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

bergson/magic/cli.py

@@ -38,7 +38,7 @@
 from ..utils.worker_utils import setup_data_pipeline
 from ..validate import load_attribution_scores, validate_scores
 from .config import MagicConfig
-from .data_stream import DataStream, pad_dataset_to_batch_size
+from .data_stream import DataStream, doc_rows, pad_dataset_to_batch_size
 from .grad_accum import accumulate_grads
 from .trainer import BackwardState, TrainerState, prepare_trainer, write_lr_history
 
@@ -156,13 +156,29 @@ 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]
+
+    rows_by_doc = doc_rows(query_dataset, num_query_docs)
+    zero_score = trim_pads(torch.zeros_like(stream.weights.detach(), device="cpu"))
+    if main and not all(rows_by_doc):
+        empty = sum(not rows for rows in rows_by_doc)
+        print(f"[per-query MAGIC] {empty} 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 rows_by_doc[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 +188,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
-        )
         qstream = DataStream(
-            one,
+            query_dataset,
             run_cfg.batch_size,
             device=device,
             input_key=run_cfg.query.prompt_column,
-            weight_shape=(n_one,),
+            rows=rows_by_doc[qi],
+            doc_id=qi,
         )
-        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,15 +225,13 @@ 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)
 
         # Free per-query state and any temp checkpoints the backward wrote.
-        del bwd_state, qgrads, qstream, one
+        del bwd_state, qgrads, qstream
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.synchronize()

bergson/magic/data_stream.py

@@ -62,6 +62,19 @@ def pad_dataset_to_batch_size(
     return dataset, num_docs, pad_count, weight_pad_count
 
 
+def doc_rows(dataset: Dataset, num_docs: int) -> list[list[int]]:
+    """For each document, the dataset rows containing any of its tokens."""
+    if "doc_ids" not in dataset.column_names:
+        return [[i] for i in range(num_docs)]
+
+    rows: list[list[int]] = [[] for _ in range(num_docs)]
+    for r, ids in enumerate(dataset["doc_ids"]):
+        for doc_id in set(ids):
+            if doc_id < num_docs:
+                rows[doc_id].append(r)
+    return rows
+
+
 class DataStream:
     def __init__(
         self,
@@ -71,13 +84,28 @@ def __init__(
         device: torch.device | str = "cpu",
         input_key: str = "text",
         weight_shape: tuple[int, ...] | None = None,
+        rows: list[int] | None = None,
+        doc_id: int | None = None,
     ):
+        """``rows`` restricts the stream to those dataset rows, cycling them to
+        fill whole batches; ``doc_id`` restricts the loss to that document's
+        tokens, which is how a single document is scored out of rows that pack
+        several (see :func:`doc_rows`)."""
         self.batch_size = batch_size
         self.dataset = dataset
         self.device = torch.device(device)
         self.input_key = input_key
         self.n = len(dataset)
-        self.num_batches = self.n // batch_size
+        self.doc_id = doc_id
+
+        self.rows = list(range(self.n)) if rows is None else list(rows)
+        if rows is not None:
+            # Cycle rather than append dead pad rows: a pad row is a real term
+            # in the loss wherever the data weights are discarded, and an
+            # all-pad batch on some rank would scale that loss down.
+            n = len(self.rows) + (-len(self.rows)) % batch_size
+            self.rows = [self.rows[i % len(self.rows)] for i in range(n)]
+        self.num_batches = len(self.rows) // batch_size
 
         # If a shape isn't provided, assume that each sequence contains one document
         if weight_shape is None:
@@ -97,11 +125,8 @@ def requires_grad(self, value: bool):
 
     def batch_rows(self, i: int) -> list[int]:
         """The current rank's dataset row indices for batch ``i``."""
-        rng = range(
-            i * self.batch_size,
-            min((i + 1) * self.batch_size, len(self.dataset)),
-        )
-        return list(rng)[self.rank :: self.world_size]
+        rows = self.rows[i * self.batch_size : (i + 1) * self.batch_size]
+        return rows[self.rank :: self.world_size]
 
     def __getitem__(self, i: int) -> dict:
         if i < 0 or i >= len(self):
@@ -114,17 +139,30 @@ def __getitem__(self, i: int) -> dict:
             labels=batch.get("labels"),
             device=self.device,
         )
+        doc_ids = batch.get("doc_ids")
+        if doc_ids is not None:
+            doc_ids = torch.tensor(doc_ids, device=self.device)
+            # doc_ids may be longer than the per-batch padded seq_len (unpacked
+            # path stores doc_ids at dataset-wide max_len); truncate to match.
+            if doc_ids.ndim == 2:
+                doc_ids = doc_ids[:, : x.shape[1]]
+
         # If the weights are 1D, we assume they correspond to documents and look for
-        # "doc_ids" in the batch to index them. If they're 2D, they correspond to tokens
+        # "doc_ids" in the batch to index them. If they're 2D, they correspond to
+        # tokens. A doc_id-restricted stream weights by row, since its rows repeat.
         if self.weights.ndim == 2:
             # Truncate to the max sequence length in the batch to avoid indexing errors
             indices = (indices, slice(None, x.shape[1]))
-        elif "doc_ids" in batch:
-            indices = torch.tensor(batch["doc_ids"], device=self.device)
-            # doc_ids may be longer than the per-batch padded seq_len (unpacked
-            # path stores doc_ids at dataset-wide max_len); truncate to match.
-            if indices.ndim == 2:
-                indices = indices[:, : x.shape[1]]
+        elif doc_ids is not None and self.doc_id is None:
+            indices = doc_ids
+
+        # Drop the other documents sharing these rows from the loss, so it is
+        # this document's own mean cross-entropy. The shift mask is the loss
+        # denominator, so it has to follow the labels.
+        if self.doc_id is not None and doc_ids is not None and doc_ids.ndim == 2:
+            y = y.where(doc_ids == self.doc_id, -100)
+            shift_loss_mask = torch.zeros_like(y, dtype=torch.bool)
+            shift_loss_mask[:, :-1] = y[:, 1:] != -100
 
         return {
             "input_ids": x,

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

@@ -21,6 +21,7 @@
 from bergson.magic import BackwardState, DataStream, Trainer
 from bergson.magic.cli import compute_per_query_magic_scores
 from bergson.magic.config import MagicConfig
+from bergson.magic.data_stream import doc_rows
 from bergson.utils.math import weighted_causal_lm_ce
 
 TINY = "trl-internal-testing/tiny-Phi3ForCausalLM"
@@ -271,3 +272,83 @@ 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_doc_rows_and_restricted_stream():
+    """Doc 0 sits inside row 0, doc 1 straddles both rows, doc 3 was dropped by
+    chunking, row 2 is padding. A stream restricted to doc 1 cycles its rows to
+    fill the batch and masks every other document's tokens out of the loss."""
+    assert doc_rows(PACKED, 4) == [[0], [0, 1], [1], []]
+    assert doc_rows(_equal_length_docs(3), 3) == [[0], [1], [2]]
+
+    stream = DataStream(PACKED, 3, input_key="input_ids", rows=[0, 1], doc_id=1)
+    assert stream.rows == [0, 1, 0]  # cycled, not a dead pad row
+    batch = stream[0]
+    assert batch["labels"].tolist() == [
+        [-100, -100, -100, 4, 5, 6],
+        [7, 8, -100, -100, -100, -100],
+        [-100, -100, -100, 4, 5, 6],
+    ]
+    # The shift mask is the loss denominator, so it must follow the masked labels
+    # rather than the row's original ones.
+    assert torch.equal(batch["shift_loss_mask"][:, :-1], batch["labels"][:, 1:] != -100)
+
+
+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))