token-xray-ml 0.1.0__py3-none-any.whl

token_xray/__init__.py

@@ -0,0 +1,6 @@
+"""Token X-ray: measured per-token gradient weights and cross-sequence leaks through real training stacks."""
+from .api import Report, xray
+from .probe import build_probe, make_batch
+
+__all__ = ["xray", "Report", "build_probe", "make_batch"]
+__version__ = "0.1.0"

token_xray/api.py

@@ -0,0 +1,66 @@
+"""
+Public API: X-ray your own training step.
+
+    from token_xray import xray
+
+    def my_step(model, sequences):
+        # build batches exactly as your trainer does, run forward + backward (no optimizer step needed)
+        ...
+
+    report = xray(my_step, lengths=[5, 9, 3, 12, 7, 4, 10, 6])
+    report.assert_token_mean()          # every target token got weight 1/N
+    print(report.table())
+
+The probe model is a tiny LLaMA whose token ids are unique tracers. `sequences` is a list of lists of
+token ids; treat them like tokenized documents. Labels follow the usual convention (labels = input ids,
+shifted inside the loss).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+import numpy as np
+
+from . import intended
+from .probe import build_probe, make_batch
+from .readout import XRay, read, recording, unit_gradients
+
+Step = Callable[[object, list[list[int]]], None]
+
+
+@dataclass
+class Report:
+    xray: XRay
+    lengths: list[int]
+
+    def deviation_from(self, want: list[list[float]]) -> float:
+        got = self.xray.normalized
+        return max(abs(g - w) for gr, wr in zip(got, want) for g, w in zip(gr, wr))
+
+    def assert_token_mean(self, tol: float = 1e-6) -> None:
+        dev = self.deviation_from(intended.token_mean(self.lengths))
+        if dev > tol:
+            raise AssertionError(
+                f"token weights deviate from a global token mean by up to {dev:.3g} (tol {tol:g}).\n" + self.table()
+            )
+
+    def assert_pure_rescaling(self, tol: float = 1e-9) -> None:
+        if self.xray.min_cos < 1 - tol:
+            raise AssertionError(f"gradient is not a pure rescaling of each token's own gradient (min cos {self.xray.min_cos:.9f})")
+
+    def table(self) -> str:
+        lines = ["seq  len   mean(w*N)   min(w*N)   max(w*N)"]
+        for i, row in enumerate(self.xray.normalized):
+            a = np.array(row)
+            lines.append(f"{i:3d} {self.lengths[i]:4d}   {a.mean():9.4f}  {a.min():9.4f}  {a.max():9.4f}")
+        return "\n".join(lines)
+
+
+def xray(step: Step, lengths: list[int], seed: int = 0) -> Report:
+    batch = make_batch(lengths)
+    unit = unit_gradients(build_probe(batch.vocab_size, seed=seed), batch)
+    model = build_probe(batch.vocab_size, seed=seed)
+    with recording(model) as rec:
+        step(model, [list(s) for s in batch.sequences])
+    return Report(xray=read(rec.grads, unit, batch), lengths=lengths)

token_xray/cli.py

@@ -0,0 +1,51 @@
+"""token-xray command line.
+
+    token-xray report                  print the frozen scenarios as a table
+    token-xray export public/xray.json write measured results for the website
+    token-xray legacy-ga out.json      measure the gradient-accumulation step with the INSTALLED
+                                       transformers (run in .venv-hf445 for the pre-fix 4.45.2 Trainer)
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(prog="token-xray")
+    sub = p.add_subparsers(dest="cmd", required=True)
+    sub.add_parser("report")
+    e = sub.add_parser("export")
+    e.add_argument("out")
+    g = sub.add_parser("legacy-ga")
+    g.add_argument("out")
+    a = p.parse_args(argv)
+
+    if a.cmd == "legacy-ga":
+        from .scenarios import legacy_ga_panel
+
+        with open(a.out, "w") as f:
+            json.dump(legacy_ga_panel(), f)
+        return 0
+
+    from .scenarios import export
+
+    data = export()
+    if a.cmd == "export":
+        with open(a.out, "w") as f:
+            json.dump(data, f, separators=(",", ":"))
+        print(f"wrote {a.out}")
+        return 0
+    print("versions:", data["versions"])
+    print("\nAggregation (max |w*N - intended|, min cosine)")
+    for pnl in data["aggregation"]["panels"]:
+        print(f"  [{pnl['kind']:9s}] {pnl['title']:52s} {pnl['maxDeviation']:.2e}  {pnl['minCos']:.6f}")
+    print("\nLeaks (max off-diagonal)")
+    for pnl in data["leak"]["panels"]:
+        print(f"  [{pnl['kind']:9s}] {pnl['title']:52s} {pnl['maxLeak']:.2e}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

token_xray/ddp.py

@@ -0,0 +1,148 @@
+"""
+Data-parallel runs on CPU (gloo): does the weight a token receives depend on how the batch was split
+across ranks?
+
+Activation gradients are local to a rank; DDP averages PARAMETER gradients in the all-reduce after
+backward. So these runs read the input-embedding parameter gradient (`readout.embedding_row_grads`)
+after the all-reduce, on rank 0, where every rank's contribution has arrived.
+
+Two launchers:
+  * `ddp_step`: torch.multiprocessing.spawn + DistributedDataParallel around the probe, with a
+    hand-written loss in one of two reductions (global token mean, per-rank mean).
+  * `hf_trainer_ddp`: the unmodified transformers Trainer under torchrun (2 CPU processes, gloo).
+
+Every run records which sequences each rank actually received, so the intended weights are computed
+from the observed split rather than an assumed one.
+"""
+from __future__ import annotations
+
+import json
+import os
+import socket
+import subprocess
+import sys
+import tempfile
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn.functional as F
+
+from .probe import TracerBatch, build_probe, make_batch
+from .readout import embedding_row_grads
+from .runners import collate
+
+REDUCTIONS = ("global_token_mean", "per_rank_mean")
+
+
+@dataclass
+class DDPResult:
+    rows: dict[str, dict[int, torch.Tensor]]  # reduction -> token id -> all-reduced embedding gradient row
+    assignment: list[list[int]]  # sequence indices seen by each rank, in order (duplicates included)
+    local_targets: list[int]  # target tokens per rank
+    world: int
+
+
+def _free_port() -> int:
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def _shard(n_seqs: int, world: int, split, sampler: bool) -> list[list[int]]:
+    if sampler:
+        from torch.utils.data import DistributedSampler
+
+        return [list(DistributedSampler(range(n_seqs), num_replicas=world, rank=r, shuffle=False, drop_last=False))
+                for r in range(world)]
+    return [list(s) for s in split]
+
+
+def _worker(rank: int, world: int, port: int, lengths: list[int], split, sampler: bool, reductions, out: str):
+    os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port))
+    torch.set_num_threads(1)
+    dist.init_process_group("gloo", rank=rank, world_size=world)
+    try:
+        batch = make_batch(lengths)
+        mine = _shard(len(lengths), world, split, sampler)[rank]
+        b = collate([{"input_ids": batch.sequences[i]} for i in mine])
+        tgt = b["labels"][:, 1:]
+        n_local = torch.tensor(float((tgt != -100).sum()), dtype=torch.float64)
+        n_all = [torch.zeros_like(n_local) for _ in range(world)]
+        dist.all_gather(n_all, n_local)
+        rows = {}
+        for red in reductions:
+            model = build_probe(batch.vocab_size)
+            ddp = torch.nn.parallel.DistributedDataParallel(model)
+            logits = ddp(input_ids=b["input_ids"], attention_mask=b["attention_mask"]).logits
+            ce = F.cross_entropy(logits[:, :-1].reshape(-1, logits.shape[-1]), tgt.reshape(-1),
+                                 ignore_index=-100, reduction="sum")
+            if red == "global_token_mean":
+                # every rank divides by the GLOBAL token count; ×world undoes DDP's average
+                loss = ce / sum(n_all) * world
+            elif red == "per_rank_mean":
+                loss = ce / n_local  # each rank's own token mean; DDP then averages the ranks
+            else:
+                raise ValueError(red)
+            loss.backward()  # DDP all-reduces (averages) the embedding gradient here
+            rows[red] = embedding_row_grads(model, range(1, batch.vocab_size))
+        if rank == 0:
+            torch.save({"rows": rows, "assignment": _shard(len(lengths), world, split, sampler),
+                        "local_targets": [int(x) for x in n_all]}, out)
+    finally:
+        dist.destroy_process_group()
+
+
+def ddp_step(
+    lengths: list[int],
+    world: int,
+    split: list[list[int]] | None = None,
+    sampler: bool = False,
+    reductions: tuple[str, ...] = REDUCTIONS,
+) -> DDPResult:
+    """One backward pass under DDP with `world` CPU processes.
+
+    split[r] lists the sequence indices given to rank r. sampler=True instead lets
+    torch.utils.data.DistributedSampler (shuffle off, drop_last off) decide, including its padding.
+    """
+    if split is None and not sampler:
+        split = [list(range(len(lengths)))] if world == 1 else None
+        if split is None:
+            raise ValueError("give a split or set sampler=True")
+    with tempfile.TemporaryDirectory() as d:
+        out = os.path.join(d, "rank0.pt")
+        mp.spawn(_worker, args=(world, _free_port(), lengths, split, sampler, tuple(reductions), out),
+                 nprocs=world, join=True)
+        r = torch.load(out)
+    return DDPResult(rows=r["rows"], assignment=r["assignment"], local_targets=r["local_targets"], world=world)
+
+
+# ── the transformers Trainer under torchrun ────────────────────────────────
+
+
+def hf_trainer_ddp(lengths: list[int], per_device_batch: int, average_tokens_across_devices: bool = True,
+                   world: int = 2, timeout: float = 300.0) -> dict:
+    """One optimizer step of the unmodified Trainer under `torchrun --nproc_per_node world` (CPU, gloo).
+    Returns the all-reduced embedding rows (rank 0) and the sequences each rank actually received."""
+    with tempfile.TemporaryDirectory() as d:
+        spec = {"lengths": lengths, "per_device_batch": per_device_batch,
+                "average_tokens_across_devices": average_tokens_across_devices, "out": d}
+        path = os.path.join(d, "spec.json")
+        with open(path, "w") as f:
+            json.dump(spec, f)
+        env = {**os.environ, "OMP_NUM_THREADS": "1", "TOKENIZERS_PARALLELISM": "false"}
+        # Static rendezvous on loopback. `--standalone` resolves the machine's host name, which on some
+        # networks is a reverse-DNS name that never answers, and the store then waits 300 s per attempt.
+        cmd = [sys.executable, "-m", "torch.distributed.run", "--nnodes=1", f"--nproc_per_node={world}",
+               "--master_addr=127.0.0.1", f"--master_port={_free_port()}",
+               "-m", "token_xray.ddp_trainer_entry", path]
+        try:
+            p = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=timeout, start_new_session=True)
+        except subprocess.TimeoutExpired as e:
+            raise RuntimeError(f"torchrun did not finish within {timeout} s; stderr tail:\n{(e.stderr or '')[-2000:]}")
+        if p.returncode != 0:
+            raise RuntimeError(f"torchrun failed ({p.returncode}):\n{p.stderr[-3000:]}")
+        r = torch.load(os.path.join(d, "rank0.pt"))
+        seen = [json.load(open(os.path.join(d, f"seen{k}.json"))) for k in range(world)]
+    return {"rows": r["rows"], "seen": seen, "world": world, "args": r["args"]}

token_xray/ddp_trainer_entry.py

@@ -0,0 +1,79 @@
+"""
+torchrun entry point: one optimizer step of the unmodified transformers Trainer on the probe, on CPU
+with the gloo backend. Launched by `ddp.hf_trainer_ddp`; not meant to be run by hand.
+
+Writes, into the spec's output directory:
+  rank0.pt      all-reduced input-embedding gradient rows, captured before the optimizer step
+  seen<r>.json  the sequence indices rank r received, in order
+"""
+from __future__ import annotations
+
+import json
+import os
+import sys
+import tempfile
+
+import torch
+from transformers import Trainer, TrainerCallback, TrainingArguments
+
+from .probe import build_probe, make_batch
+from .readout import embedding_row_grads
+from .runners import _ListDataset, collate
+
+
+def main(path: str) -> None:
+    spec = json.load(open(path))
+    rank = int(os.environ.get("RANK", "0"))
+    batch = make_batch(spec["lengths"])
+    model = build_probe(batch.vocab_size)
+    ds = _ListDataset(batch)
+    seen: list[int] = []
+
+    class Capture(TrainerCallback):
+        def on_pre_optimizer_step(self, args, state, control, **kw):
+            # DDP has all-reduced (averaged) the parameter gradients during the last backward
+            if rank == 0:
+                rows = embedding_row_grads(model, range(1, batch.vocab_size))
+                torch.save({"rows": rows, "args": {"average_tokens_across_devices": args.average_tokens_across_devices,
+                                                   "world_size": args.world_size}},
+                           os.path.join(spec["out"], "rank0.pt"))
+
+    with tempfile.TemporaryDirectory() as out:
+        args = TrainingArguments(
+            output_dir=out,
+            per_device_train_batch_size=spec["per_device_batch"],
+            gradient_accumulation_steps=1,
+            max_steps=1,
+            learning_rate=1e-3,
+            max_grad_norm=0.0,
+            report_to=[],
+            use_cpu=True,
+            ddp_backend="gloo",
+            save_strategy="no",
+            logging_strategy="no",
+            seed=0,
+            data_seed=0,
+            remove_unused_columns=False,
+            dataloader_drop_last=False,
+            disable_tqdm=True,
+            average_tokens_across_devices=spec["average_tokens_across_devices"],
+        )
+        trainer = Trainer(model=model, args=args, train_dataset=ds, data_collator=lambda f: collate(f),
+                          callbacks=[Capture()])
+        trainer._get_train_sampler = lambda *a, **k: torch.utils.data.SequentialSampler(ds)
+        compute = trainer.compute_loss
+
+        def spy(m, inputs, *a, **k):
+            for row in inputs["input_ids"].tolist():
+                first = next(t for t in row if t != 0)
+                seen.append(batch.seq_of_token[first])
+            return compute(m, inputs, *a, **k)
+
+        trainer.compute_loss = spy
+        trainer.train()
+    with open(os.path.join(spec["out"], f"seen{rank}.json"), "w") as f:
+        json.dump(seen, f)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1])

token_xray/intended.py

@@ -0,0 +1,162 @@
+"""
+Intended weights: what the equation says each token should receive (closed form).
+
+All functions return per-sequence lists of weights for target tokens (length L_i - 1 each),
+already multiplied by N (total target tokens), so a correct global token mean is 1.0 everywhere.
+"""
+from __future__ import annotations
+
+import math
+
+
+def _targets(lengths: list[int]) -> list[int]:
+    return [L - 1 for L in lengths]
+
+
+def token_mean(lengths: list[int]) -> list[list[float]]:
+    """Global mean over all target tokens in the effective batch: w = 1/N."""
+    return [[1.0] * n for n in _targets(lengths)]
+
+
+def sequence_mean(lengths: list[int]) -> list[list[float]]:
+    """Mean over tokens within each sequence, then over sequences: w = 1/(B * n_i)."""
+    n = _targets(lengths)
+    N, B = sum(n), len(n)
+    return [[N / (B * ni)] * ni for ni in n]
+
+
+def mean_of_microbatch_means(lengths: list[int], micro: int) -> list[list[float]]:
+    """Token mean inside each micro-batch, then averaged over micro-batches (sequential split)."""
+    n = _targets(lengths)
+    N = sum(n)
+    groups = [n[i : i + micro] for i in range(0, len(n), micro)]
+    k = len(groups)
+    out = []
+    for g in groups:
+        Ng = sum(g)
+        out += [[N / (k * Ng)] * ni for ni in g]
+    return out
+
+
+def first_microbatch_only(lengths: list[int], micro: int) -> list[list[float]]:
+    """Only the first micro-batch survives; its tokens keep the global normalizer."""
+    n = _targets(lengths)
+    return [[1.0 if i < micro else 0.0] * ni for i, ni in enumerate(n)]
+
+
+
+def labels_equal_to_pad_masked(lengths: list[int]) -> list[list[float]]:
+    """Each sequence ends with an EOS whose id is also the pad id. A collator that sets every label equal
+    to the pad id to -100 drops each sequence's last target (the EOS): it gets weight 0, and the
+    remaining tokens share the normaliser, w·N = N / (N - B). Lengths include the EOS."""
+    n = _targets(lengths)
+    N, B = sum(n), len(n)
+    return [[N / (N - B)] * (ni - 1) + [0.0] for ni in n]
+
+# ── preference and RL objectives ───────────────────────────────────────────
+#
+# These return the raw weight on each completion token's cross-entropy, w = −∂L/∂log π(token)
+# (positive = the token is reinforced), not multiplied by N. Rows are completions, in batch order.
+
+
+def _sigmoid(x: float) -> float:
+    return 1.0 / (1.0 + math.exp(-x)) if x >= 0 else math.exp(x) / (1.0 + math.exp(x))
+
+
+def dpo_sigmoid(
+    chosen_lens: list[int], rejected_lens: list[int], deltas: list[float], beta: float
+) -> tuple[list[list[float]], list[list[float]]]:
+    """DPO (Rafailov et al., 2023), sigmoid loss averaged over P pairs:
+
+        L = −1/P Σ_i log σ(β Δ_i),   Δ_i = (log π_c − log π_ref,c) − (log π_r − log π_ref,r)
+
+    with sequence log-probs summed over completion tokens, so every token of a completion gets the
+    same weight: +β σ(−β Δ_i) / P on the chosen tokens, −β σ(−β Δ_i) / P on the rejected ones.
+    """
+    P = len(deltas)
+    if not (len(chosen_lens) == len(rejected_lens) == P):
+        raise ValueError("one chosen length, one rejected length and one margin per pair")
+    k = [beta * _sigmoid(-beta * d) / P for d in deltas]
+    return [[k[i]] * n for i, n in enumerate(chosen_lens)], [[-k[i]] * n for i, n in enumerate(rejected_lens)]
+
+
+def group_advantages(rewards: list[float], scale: bool = True, eps: float = 1e-4) -> list[float]:
+    """GRPO advantages for one group: (r − mean) / (std + eps), std with ddof = 1 (torch default)."""
+    G = len(rewards)
+    m = sum(rewards) / G
+    c = [r - m for r in rewards]
+    if not scale:
+        return c
+    sd = math.sqrt(sum(x * x for x in c) / (G - 1))
+    return [x / (sd + eps) for x in c]
+
+
+def group_policy_gradient(
+    lengths: list[int],
+    rewards: list[float],
+    loss_type: str,
+    ratios: list[list[float]] | None = None,
+    epsilon_low: float = 0.2,
+    epsilon_high: float | None = None,
+    max_completion_length: int | None = None,
+    scale_rewards: bool = True,
+) -> tuple[list[float], list[list[float]]]:
+    """One group of G completions, one optimizer step, KL off (beta = 0), clipped surrogate
+
+        ℓ_t = −min(ρ_t A_i, clip(ρ_t, 1 − ε_low, 1 + ε_high) A_i)
+
+    reduced as the installed TRL does for each loss_type (single process, no accumulation):
+        grpo    : mean over completions of the per-completion token mean   (norm G · |o_i|)
+        bnpo    : token mean over the batch                                (norm Σ_j |o_j|)
+        dapo    : token mean over the generation batch                     (norm Σ_j |o_j|)
+        dr_grpo : sum over tokens / (G · max_completion_length)            (norm G · L_max)
+
+    Weight on token t of completion i: A_i ρ_t / norm, or 0 where the clipped branch is selected
+    (A > 0 and ρ > 1 + ε_high, or A < 0 and ρ < 1 − ε_low). Returns (advantages, weights).
+    """
+    G = len(lengths)
+    eh = epsilon_low if epsilon_high is None else epsilon_high
+    A = group_advantages(rewards, scale=scale_rewards)
+    total = sum(lengths)
+    rho = ratios if ratios is not None else [[1.0] * L for L in lengths]
+    out = []
+    for i, L in enumerate(lengths):
+        if loss_type == "grpo":
+            norm = G * L
+        elif loss_type in ("bnpo", "dapo"):
+            norm = total
+        elif loss_type == "dr_grpo":
+            if max_completion_length is None:
+                raise ValueError("dr_grpo normalises by max_completion_length")
+            norm = G * max_completion_length
+        else:
+            raise ValueError(f"unsupported loss_type {loss_type!r}")
+        row = []
+        for r in rho[i]:
+            clipped = (A[i] > 0 and r > 1 + eh) or (A[i] < 0 and r < 1 - epsilon_low)
+            row.append(0.0 if clipped else A[i] * r / norm)
+        out.append(row)
+    return A, out
+
+
+# ── data parallel ──────────────────────────────────────────────────────────
+
+
+def per_rank_mean(lengths: list[int], rank_of: list[int], world: int) -> list[list[float]]:
+    """Token mean on each rank, then DDP's average over ranks: w = 1 / (W · N_r), times N.
+    Equal to the global token mean only when every rank holds the same number of target tokens."""
+    n = _targets(lengths)
+    N = sum(n)
+    per = [0] * world
+    for ni, r in zip(n, rank_of):
+        per[r] += ni
+    return [[N / (world * per[r])] * ni for ni, r in zip(n, rank_of)]
+
+
+def with_multiplicity(lengths: list[int], copies: list[int]) -> list[list[float]]:
+    """Global token mean over a batch in which sequence i appears copies[i] times (sampler padding):
+    each copy counts in the normaliser and each adds its gradient, so w·N = copies_i · N / N_seen."""
+    n = _targets(lengths)
+    N = sum(n)
+    seen = sum(c * ni for c, ni in zip(copies, n))
+    return [[c * N / seen] * ni for c, ni in zip(copies, n)]

token_xray/leak.py

@@ -0,0 +1,75 @@
+"""
+Leak mode: does one sequence's loss send gradient into another sequence?
+
+With attention live, run K backward passes. In pass i, only sequence i keeps its labels. The
+gradient that reaches the positions of any other sequence j is a leak. On a correct stack the leak
+matrix is exactly zero off the diagonal.
+
+Layouts:
+  * padded:            one sequence per row, right padding, attention mask. Reference: no leaks.
+  * packed_naive:      sequences concatenated into one row with no boundary information.
+  * packed_position_ids: concatenated, position_ids restart at 0 for each sequence (the
+                       padding-free / flattening convention), no attention mask.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from .probe import TracerBatch, build_probe
+from .readout import recording
+from .runners import collate
+
+LAYOUTS = ("padded", "packed_naive", "packed_position_ids")
+
+
+def _inputs(batch: TracerBatch, layout: str, keep: int):
+    seqs = batch.sequences
+    if layout == "padded":
+        b = collate([{"input_ids": s} for s in seqs])
+        lab = b["labels"].clone()
+        for i in range(len(seqs)):
+            if i != keep:
+                lab[i] = -100
+        return {"input_ids": b["input_ids"], "attention_mask": b["attention_mask"]}, lab
+    ids = torch.tensor([[t for s in seqs for t in s]])
+    lab = torch.full_like(ids, -100)
+    start = sum(len(s) for s in seqs[:keep])
+    L = len(seqs[keep])
+    lab[0, start : start + L] = ids[0, start : start + L]
+    # the first token of a sequence must never be a target of the previous sequence's last token
+    lab[0, start] = -100
+    inputs = {"input_ids": ids}
+    if layout == "packed_position_ids":
+        inputs["position_ids"] = torch.tensor([[p for s in seqs for p in range(len(s))]])
+    return inputs, lab
+
+
+def leak_matrix(batch: TracerBatch, layout: str, attn_implementation: str = "sdpa", seed: int = 0) -> np.ndarray:
+    """M[i, j] = ||grad reaching sequence j from sequence i's loss|| / ||grad reaching sequence i||."""
+    if layout not in LAYOUTS:
+        raise ValueError(f"layout must be one of {LAYOUTS}")
+    K = len(batch.sequences)
+    M = np.zeros((K, K))
+    for i in range(K):
+        model = build_probe(batch.vocab_size, "leak", seed=seed, attn_implementation=attn_implementation)
+        inputs, lab = _inputs(batch, layout, i)
+        with recording(model) as rec:
+            logits = model(**inputs).logits
+            loss = F.cross_entropy(
+                logits[:, :-1].reshape(-1, logits.shape[-1]), lab[:, 1:].reshape(-1), ignore_index=-100, reduction="sum"
+            )
+            loss.backward()
+        norms = np.zeros(K)
+        for tok, g in rec.grads.items():
+            norms[batch.seq_of_token[tok]] += float(g.norm()) ** 2
+        norms = np.sqrt(norms)
+        M[i] = norms / norms[i] if norms[i] > 0 else norms
+    return M
+
+
+def max_leak(M: np.ndarray) -> float:
+    off = M.copy()
+    np.fill_diagonal(off, 0.0)
+    return float(off.max())