From 829e77ee55efd1f91c4dc1f163a172349a1beebb Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 21:33:26 +0900 Subject: [PATCH 1/2] fix(magic): score per-query MAGIC by document, not by row MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. `compute_per_query_magic_scores` selected row `qi` for query `qi` and sized the stream's weights by row, so `DataStream` — which indexes the weights by document id whenever the batch has a `doc_ids` column — walked off the end: IndexError: index 2 is out of bounds for dimension 0 with size 2 bergson/magic/data_stream.py:132 self.weights[indices] Repro: `bergson magic --query.chunk_length 32 ...`. Teach `DataStream` about documents instead. It already owns the `doc_ids` convention, so `doc_rows()` lives beside it, and two optional arguments make a stream over one document: `rows` restricts it to those dataset rows and `doc_id` masks every other document's tokens out of the labels. Each score column is then that document's own mean cross-entropy, the unit `validate_scores` and `per_doc_query_losses` already score against. `shift_loss_mask` is rebuilt with the labels since it is the loss denominator. `rows` cycles through the document's rows rather than appending dead pad ones — a pad row is a real term in the query loss, and an all-pad batch on some rank would silently scale the query gradient down. With one row per document this repeats that row, exactly as `pad_dataset_to_batch_size` did, so unchunked runs are unchanged. Chunking drops the tail that doesn't fill a chunk, so a short document can leave no tokens behind; those score zeros, matching the zero baseline loss `per_doc_query_losses` gives them. Their correlation is undefined, so the reported mean Spearman now averages the queries that have one and reports the nan count. Co-Authored-By: Claude Opus 5 (1M context) --- bergson/magic/cli.py | 35 +++++++++------ bergson/magic/data_stream.py | 71 ++++++++++++++++++++++++------ bergson/validate.py | 7 ++- tests/test_per_query_magic.py | 81 +++++++++++++++++++++++++++++++++++ 4 files changed, 167 insertions(+), 27 deletions(-) diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index 47434c78..c183539a 100644 --- a/bergson/magic/cli.py +++ b/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,6 +156,17 @@ 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") @@ -163,6 +174,11 @@ def compute_per_query_magic_scores( 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() diff --git a/bergson/magic/data_stream.py b/bergson/magic/data_stream.py index fb031879..3c659fbd 100644 --- a/bergson/magic/data_stream.py +++ b/bergson/magic/data_stream.py @@ -62,6 +62,26 @@ 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]]: + """Row indices holding each document's tokens. + + A chunked dataset (``chunk_length > 0``) packs several documents into a row + and splits documents across rows, so document ``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 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 +91,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 +132,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 +146,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, diff --git a/bergson/validate.py b/bergson/validate.py index 71c3805f..c182a7cb 100644 --- a/bergson/validate.py +++ b/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( diff --git a/tests/test_per_query_magic.py b/tests/test_per_query_magic.py index a8910f8e..c7fabdf9 100644 --- a/tests/test_per_query_magic.py +++ b/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)) From fcbebca7eea85960d39880c8523e43411ef949d1 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 22:06:45 +0900 Subject: [PATCH 2/2] docs: trim doc_rows docstring to its summary line The body restated the chunking behaviour the function's own body shows, and led with "Row indices holding each document's tokens", which reads as a document-to-row lookup without saying which direction or what a row is. Co-Authored-By: Claude Opus 5 (1M context) --- bergson/magic/data_stream.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/bergson/magic/data_stream.py b/bergson/magic/data_stream.py index 3c659fbd..7946d06f 100644 --- a/bergson/magic/data_stream.py +++ b/bergson/magic/data_stream.py @@ -63,14 +63,7 @@ def pad_dataset_to_batch_size( def doc_rows(dataset: Dataset, num_docs: int) -> list[list[int]]: - """Row indices holding each document's tokens. - - A chunked dataset (``chunk_length > 0``) packs several documents into a row - and splits documents across rows, so document ``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. - """ + """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)]