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

bergson/magic/cli.py

@@ -110,6 +110,52 @@ def compute_query_gradients(
     return grad_accum, float(loss_accum)
 
 
+def query_doc_rows(query_dataset: Dataset, num_query_docs: int) -> list[list[int]]:
+    """Row indices holding each query document's tokens.
+
+    A chunked query set (``query.chunk_length > 0``) packs several documents
+    into a row and splits documents across rows, so query ``i`` is not row
+    ``i``; its per-token ``doc_ids`` column says which is which. Pad rows carry
+    a synthetic id past the real documents, and a document that chunking
+    dropped (the tail that doesn't fill a chunk) gets no rows.
+    """
+    if "doc_ids" not in query_dataset.column_names:
+        return [[i] for i in range(num_query_docs)]
+
+    rows: list[list[int]] = [[] for _ in range(num_query_docs)]
+    for r, ids in enumerate(query_dataset["doc_ids"]):
+        for doc_id in set(ids):
+            if doc_id < num_query_docs:
+                rows[doc_id].append(r)
+    return rows
+
+
+def query_doc_batch(
+    query_dataset: Dataset, doc_id: int, rows: list[int], batch_size: int
+) -> Dataset:
+    """One document's rows, other documents' tokens masked out with ``-100``,
+    padded to a whole batch by cycling through those rows.
+
+    The loss is then this document's own mean cross-entropy. Dropping
+    ``doc_ids`` keeps ``DataStream`` indexing the weights by row; cycling
+    rather than appending dead pad rows keeps every row a real term in the
+    loss, since ``compute_query_gradients`` discards the weights that would
+    silence one.
+    """
+    total = len(rows) + (-len(rows)) % batch_size
+    one = query_dataset.select([rows[i % len(rows)] for i in range(total)])
+    if "doc_ids" not in one.column_names:
+        return one
+
+    label_col = "labels" if "labels" in one.column_names else "input_ids"
+    labels = [
+        [tok if d == doc_id else -100 for tok, d in zip(toks, ids)]
+        for toks, ids in zip(one[label_col], one["doc_ids"])
+    ]
+    drop = [c for c in ("doc_ids", "labels") if c in one.column_names]
+    return one.remove_columns(drop).add_column("labels", labels)
+
+
 def compute_per_query_magic_scores(
     trainer,
     ckpts_path: str,
@@ -156,13 +202,30 @@ def compute_per_query_magic_scores(
         if isinstance(buf, torch.Tensor) and buf.is_floating_point()
     ]
 
+    def trim_pads(s: torch.Tensor) -> torch.Tensor:
+        if not pad_count:
+            return s
+        return s[:-weight_pad_count] if s.ndim == 1 else s[:-pad_count]
+
+    doc_rows = query_doc_rows(query_dataset, num_query_docs)
+    zero_score = trim_pads(torch.zeros_like(stream.weights.detach(), device="cpu"))
+    if main and not all(doc_rows):
+        print(
+            f"[per-query MAGIC] {sum(not r for r in doc_rows)} queries have no tokens"
+        )
+
     per_query = []
     for qi in range(num_query_docs):
         qpath = os.path.join(scores_dir, f"q{qi}.pt")
         if os.path.exists(qpath):  # resume: already scored
             per_query.append(torch.load(qpath, map_location="cpu"))
             continue
 
+        if not doc_rows[qi]:
+            # No tokens, no gradient — and validate_scores baselines it at 0 too.
+            per_query.append(zero_score)
+            continue
+
         # Restore the final trained state (the backward walks it back down the
         # trajectory). detach_ first: the previous iteration left params
         # requiring grad, and copy_ is an in-place write a leaf-requiring-grad
@@ -172,19 +235,14 @@ def compute_per_query_magic_scores(
         fwd_state.copy_(restored)
         del restored
 
-        one = query_dataset.select([qi])
-        one, n_one, one_pad, one_wpad = pad_dataset_to_batch_size(
-            one, run_cfg.batch_size, 1, f"Query {qi}", global_rank
-        )
+        one = query_doc_batch(query_dataset, qi, doc_rows[qi], run_cfg.batch_size)
         qstream = DataStream(
             one,
             run_cfg.batch_size,
             device=device,
             input_key=run_cfg.query.prompt_column,
-            weight_shape=(n_one,),
+            weight_shape=(len(one),),
         )
-        if one_pad:
-            qstream.weights.data[-one_wpad:] = 0.0
         qgrads, _ = compute_query_gradients(
             fwd_state, model, qstream, "mean", run_cfg.fsdp, run_cfg.grad_accum_steps
         )
@@ -214,9 +272,7 @@ def compute_per_query_magic_scores(
         if world_size > 1:
             dist.all_reduce(bwd_state.weight_grads, op=dist.ReduceOp.SUM)
 
-        s = bwd_state.weight_grads.detach().cpu()
-        if pad_count:
-            s = s[:-weight_pad_count] if s.ndim == 1 else s[:-pad_count]
+        s = trim_pads(bwd_state.weight_grads.detach().cpu())
         if main:
             torch.save(s, qpath)
         per_query.append(s)

bergson/validate.py

@@ -228,7 +228,12 @@ def report_multi_query_validation(
             float(baselines[q]),
         )
     summary_csv_writer.close()
-    print(f"Mean Spearman across {num_queries} queries: {np.mean(rhos):.4f}")
+
+    # A query whose loss doesn't move across subsets (an empty document that
+    # chunking left with no tokens, say) has no correlation to average in.
+    rhos = [rho for rho in rhos if not np.isnan(rho)]
+    mean = np.mean(rhos) if rhos else float("nan")
+    print(f"Mean Spearman {mean:.4f} ({num_queries - len(rhos)}/{num_queries} nan)")
 
 
 def validate_scores(

tests/test_per_query_magic.py

@@ -19,7 +19,11 @@
 
 from bergson.distributed import grad_tree
 from bergson.magic import BackwardState, DataStream, Trainer
-from bergson.magic.cli import compute_per_query_magic_scores
+from bergson.magic.cli import (
+    compute_per_query_magic_scores,
+    query_doc_batch,
+    query_doc_rows,
+)
 from bergson.magic.config import MagicConfig
 from bergson.utils.math import weighted_causal_lm_ce
 
@@ -271,3 +275,80 @@ 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
+
+
+# A chunked query set (query.chunk_length > 0) 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.
+
+PACKED = Dataset.from_dict(
+    {
+        "input_ids": [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12], [0] * 6],
+        "doc_ids": [[0, 0, 0, 1, 1, 1], [1, 1, 2, 2, 2, 2], [4] * 6],
+    }
+)
+
+
+def test_query_doc_rows_and_batch():
+    """Doc 0 sits inside row 0, doc 1 straddles both rows, doc 3 was dropped by
+    chunking, row 2 is padding. Doc 1's batch keeps only its tokens, drops
+    doc_ids, and pads by cycling through its own rows."""
+    assert query_doc_rows(PACKED, 4) == [[0], [0, 1], [1], []]
+    assert query_doc_rows(_equal_length_docs(3), 3) == [[0], [1], [2]]
+
+    one = query_doc_batch(PACKED, 1, [0, 1], batch_size=3)
+    assert "doc_ids" not in one.column_names
+    assert one["input_ids"][2] == one["input_ids"][0]  # cycled, not a dead pad
+    assert one["labels"] == [
+        [-100, -100, -100, 4, 5, 6],
+        [7, 8, -100, -100, -100, -100],
+        [-100, -100, -100, 4, 5, 6],
+    ]
+
+
+def _per_query(query_ds, num_query_docs):
+    model = _model()
+    opt = torchopt.adamw(1e-4, betas=(0.95, 0.975), eps_root=1e-2)
+    trainer, fwd_state = Trainer.initialize(model, opt)
+    stream = DataStream(_equal_length_docs(3), batch_size=1, device="cpu")
+
+    with tempfile.TemporaryDirectory() as run_path:
+        ckpts = f"{run_path}/checkpoints"
+        fwd_state = trainer.train(fwd_state, stream, inplace=True, save_dir=ckpts)
+        cfg = MagicConfig(run_path=run_path, query_method="none", batch_size=2)
+        cfg.query.prompt_column = "input_ids"
+        stream.requires_grad = True
+        return compute_per_query_magic_scores(
+            trainer,
+            ckpts,
+            stream,
+            fwd_state,
+            model,
+            query_ds,
+            num_query_docs,
+            cfg,
+            1,
+            0,
+            0,
+            0,
+        )
+
+
+def test_per_query_packed_docs_score_separately():
+    """Two documents packed into one row (which used to crash: the stream sizes
+    its weights by row, DataStream indexes them by doc id) score as the same two
+    documents split across rows — each column is that document's tokens alone.
+    Doc 2 is missing from the rows entirely, as chunking's dropped tail is."""
+    row = list(range(100, 106))
+    packed = Dataset.from_dict({"input_ids": [row], "doc_ids": [[0, 0, 0, 1, 1, 1]]})
+    split = Dataset.from_dict(
+        {
+            "input_ids": [row, row],
+            "labels": [row[:3] + [-100] * 3, [-100] * 3 + row[3:]],
+        }
+    )
+
+    scores = _per_query(packed, 3)
+    assert not torch.allclose(scores[:, 0], scores[:, 1])
+    assert (scores[:, 2] == 0).all()
+    torch.testing.assert_close(scores[:, :2], _per_query(split, 2))