diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index 47434c78..55257ba9 100644 --- a/bergson/magic/cli.py +++ b/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,6 +202,18 @@ 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") @@ -163,6 +221,11 @@ def compute_per_query_magic_scores( 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) 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..a51b457b 100644 --- a/tests/test_per_query_magic.py +++ b/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))