From 22d2dee0cba939367779fa0b61f88913475be153 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Fri, 17 Jul 2026 09:18:16 +0000 Subject: [PATCH 01/10] feat: ADAM SOURCE and per-checkpoint optimizer saving in Trainer --- .../approx_unrolling/adam_preconditioner.py | 193 +++++++++++ .../approx_unrolling/approx_unrolling_math.py | 24 +- bergson/approx_unrolling/pipeline.py | 44 +++ bergson/config/config.py | 11 +- bergson/hessians/apply_hessian.py | 48 ++- bergson/hessians/preconditioner.py | 175 +++++++++- bergson/magic/cli.py | 6 +- tests/test_distributed_magic.py | 20 ++ tests/test_source_optimizer_variants.py | 312 ++++++++++++++++++ 9 files changed, 813 insertions(+), 20 deletions(-) create mode 100644 bergson/approx_unrolling/adam_preconditioner.py create mode 100644 tests/test_source_optimizer_variants.py diff --git a/bergson/approx_unrolling/adam_preconditioner.py b/bergson/approx_unrolling/adam_preconditioner.py new file mode 100644 index 00000000..30ca2187 --- /dev/null +++ b/bergson/approx_unrolling/adam_preconditioner.py @@ -0,0 +1,193 @@ +"""Per-segment Adam/AdamW diagonal preconditioners for the preconditioned +SOURCE (approx-unrolling) variant. + +Loads optimizer.pt from the checkpoint dirs. You can use +``TrainingConfig.save_optimizer_state`` to save them with the Bergson +Trainer. They contain the second moment buffer, ``step``, and ``betas``). +:func:`build_segment_preconditioners` bias-corrects the moments +(v_hat = exp_avg_sq / (1 - beta2^step)), matches them to the segment's +EKFAC factor modules, orients them to bergson's ``[out, in]`` gradient +grids (params stored transposed, e.g. GPT-2 Conv1D, are flipped), averages +within each segment (the same stationary-within-segment averaging used for +the K-FAC factors), and writes the diagonal preconditioner + + P = 1 / (sqrt(v_hat_bar + eps_root) + eps) + +to ``/segment_/preconditioner.safetensors``. +""" + +from glob import glob +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file +from transformers import AutoConfig, AutoModelForCausalLM + +from ..utils.load_from_optimizer import ( + get_unfactored_second_moment, + load_optimizer, + optimizer_param_index_to_name, + orient_second_moment, +) + +OPTIMIZER_STATE_FILE = "optimizer.pt" + + +def _module_grids(hessian_dir: str | Path) -> dict[str, tuple[int, int]]: + """Module name -> full [out, in] grid shape, read from the factor shards. + + The activation eigenvectors are ``[c_i, I]`` (columns = full in-dim) and + the gradient eigenvectors ``[c_o, O]``, so shard_0's column counts give + the full grid shape regardless of world size. + """ + grids: dict[str, tuple[int, int]] = {} + a_dir = Path(hessian_dir) / "eigen_activation_sharded" + g_dir = Path(hessian_dir) / "eigen_gradient_sharded" + if not a_dir.is_dir() or not glob(str(g_dir / "shard_*.safetensors")): + raise FileNotFoundError(f"No EKFAC factor shards under {hessian_dir}") + with ( + safe_open(str(a_dir / "shard_0.safetensors"), framework="pt") as fa, + safe_open(str(g_dir / "shard_0.safetensors"), framework="pt") as fg, + ): + for name in fa.keys(): + in_dim = fa.get_slice(name).get_shape()[1] + out_dim = fg.get_slice(name).get_shape()[1] + grids[name] = (out_dim, in_dim) + return grids + + +def _match_params( + grids: dict[str, tuple[int, int]], param_names: list[str] +) -> dict[str, str]: + """Factor module name -> param name. Factor names are suffixes of the + model's param names (e.g. ``h.0.attn.c_attn`` -> + ``transformer.h.0.attn.c_attn.weight``).""" + matches: dict[str, str] = {} + for module in grids: + suffix = f"{module}.weight" + candidates = [n for n in param_names if n == suffix or n.endswith(f".{suffix}")] + if len(candidates) != 1: + raise KeyError( + f"Module {module!r} matched {len(candidates)} params " + f"({candidates}); expected exactly one *.{suffix}." + ) + matches[module] = candidates[0] + return matches + + +def _orient(p: torch.Tensor, grid: tuple[int, int], module: str, layer) -> torch.Tensor: + """Orient a param-shaped grid to bergson's [out, in] layout, raising on a + shape that fits neither orientation.""" + oriented = orient_second_moment(p, *grid, layer=layer) + if oriented is None: + raise ValueError( + f"Second moments for {module!r} have shape {tuple(p.shape)}; " + f"expected {grid} or its transpose." + ) + return oriented + + +def _adam_hparams(optimizer_pt: dict) -> tuple[float, float, float]: + """(beta2, eps, eps_root) from a standard ``optimizer.pt``'s param_groups. + + ``betas`` and ``eps`` are standard AdamW param-group keys (HF Trainer + checkpoints carry them); ``eps_root`` is torchopt-specific and defaults + to 0 (standard AdamW has no epsilon inside the sqrt).""" + groups = optimizer_pt.get("param_groups", []) + betas = groups[0].get("betas") if groups else None + eps = groups[0].get("eps") if groups else None + if betas is None or eps is None: + raise ValueError( + "optimizer.pt has no param_groups betas/eps; regenerate the " + "snapshot exports with TrainingConfig.save_optimizer_state " + "(older exports lack the hyperparameters needed to build the " + "preconditioner)." + ) + return float(betas[1]), float(eps), float(groups[0].get("eps_root", 0.0)) + + +def _extract_v_hat( + optimizer_pt: dict, index_to_name: dict[int, str], beta2: float +) -> dict[str, torch.Tensor]: + """Bias-corrected second moments keyed by param name from a standard + ``optimizer.pt`` dict carrying a per-entry ``step``.""" + + v_hat: dict[str, torch.Tensor] = {} + for idx, entry in optimizer_pt["state"].items(): + idx = int(idx) + # Prefer the recorded param_name: positional indices written from an + # FSDP-wrapped model do not match a fresh model's enumeration. + name = entry.get("param_name") or index_to_name.get(idx) + if name is None: + continue + if "step" not in entry: + raise ValueError( + f"optimizer.pt state entry {idx} has no step; regenerate the " + "snapshot exports with TrainingConfig.save_optimizer_state." + ) + step = int( + entry["step"].item() if torch.is_tensor(entry["step"]) else entry["step"] + ) + nu = get_unfactored_second_moment(entry).to(torch.float32) + v_hat[name] = nu / (1.0 - beta2**step) if step > 0 else torch.zeros_like(nu) + return v_hat + + +def build_segment_preconditioners( + run_path: str | Path, + method: str, + checkpoints: list[str], + segments: int, +) -> list[Path]: + """Build ``/segment_/preconditioner.safetensors`` for every + segment from the checkpoints' ``optimizer.pt`` files; see module docs. + The Adam hyperparameters (betas/eps/eps_root) are read from the files. + + Returns the per-segment preconditioner paths. + """ + base = Path(run_path) + per_segment = len(checkpoints) // segments + # For the optimizer.pt index mapping (used when entries carry no + # param_name, e.g. HF Trainer checkpoints). + config = AutoConfig.from_pretrained(checkpoints[0]) + reference_model = AutoModelForCausalLM.from_config(config) + out_paths: list[Path] = [] + for seg in range(segments): + seg_ckpts = checkpoints[seg * per_segment : (seg + 1) * per_segment] + v_hat_sum: dict[str, torch.Tensor] = {} + eps = eps_root = 0.0 + for ckpt in seg_ckpts: + if not (Path(ckpt) / OPTIMIZER_STATE_FILE).exists(): + raise FileNotFoundError( + f"use_adam_preconditioner requires " + f"{Path(ckpt) / OPTIMIZER_STATE_FILE}; write it during " + "training with TrainingConfig.save_optimizer_state and " + "copy it into the exported checkpoint dir." + ) + optimizer_pt = load_optimizer(str(ckpt)) + beta2, eps, eps_root = _adam_hparams(optimizer_pt) + index_to_name = optimizer_param_index_to_name(optimizer_pt, reference_model) + for name, v_hat in _extract_v_hat( + optimizer_pt, index_to_name, beta2 + ).items(): + v_hat_sum[name] = v_hat_sum.get(name, 0) + v_hat + + grids = _module_grids(base / f"segment_{seg}" / method) + matches = _match_params(grids, list(v_hat_sum)) + precond = {} + for module, param_name in matches.items(): + try: + layer = reference_model.get_submodule( + param_name.removesuffix(".weight") + ) + except AttributeError: + layer = None + v_hat_bar = v_hat_sum[param_name] / len(seg_ckpts) + p = 1.0 / ((v_hat_bar + eps_root).sqrt() + eps) + precond[module] = _orient(p, grids[module], module, layer) + + out_path = base / f"segment_{seg}" / "preconditioner.safetensors" + save_file(precond, str(out_path)) + out_paths.append(out_path) + return out_paths diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index bf093272..a96e9e30 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -60,6 +60,7 @@ def compute_lr_times_steps_per_segment( ) -> list[float]: """Per-segment lr * K. Use ``lr_list * step_size_list`` if set on config; else equal-partition log_history.json into segments and sum per-step LRs. +<<<<<<< HEAD With SGD heavy-ball ``momentum`` beta, scale by the terminal velocity 1/(1-beta) (Bae et al. 2024, App. D.2). @@ -71,6 +72,15 @@ def compute_lr_times_steps_per_segment( raise ValueError(f"momentum must be in [0, 1), got {momentum}.") momentum_scale = 1.0 / (1.0 - momentum) +======= + With SGD heavy-ball ``momentum`` beta, scale by the terminal velocity + 1/(1-beta) (Bae et al., 2024, Appendix D.2).""" + cfg = approx_unrolling_cfg + L = cfg.segments + if not 0.0 <= cfg.momentum < 1.0: + raise ValueError(f"momentum must be in [0, 1), got {cfg.momentum}.") + momentum_scale = 1.0 / (1.0 - cfg.momentum) +>>>>>>> 45df0de8 (feat: ADAM SOURCE and per-checkpoint optimizer saving in Trainer) if cfg.lr_list and cfg.step_size_list: return [ lr * k * momentum_scale for lr, k in zip(cfg.lr_list, cfg.step_size_list) @@ -136,18 +146,24 @@ def apply_eigfn_to_query( lr_times_steps: float, fn_kind: str, distributed: DistributedConfig, + precond_path: str = "", ) -> None: """Apply F_segment or F_backward of one segment to a stored query gradient. ``fn_kind`` is "f_segment" or "f_backward". The segment eigenvalues are already checkpoint-averaged (expected eigenvalues), so the eigenfunction is - applied to them directly. - """ + applied to them directly. ``precond_path`` selects the + preconditioned-optimizer variant: the eigenfunction is evaluated on a + diagonal approximation of P^1/2 H P^1/2 in parameter space, with + F_segment's output additionally multiplied by P (its P^1/2 . P^1/2 + sandwich; F_backward's sandwich cancels).""" cfg = EkfacConfig( hessian_method_path=str(segment_dir), gradient_path=str(src_grad_path), run_path=str(dst_grad_path), ev_correction=True, + precond_path=precond_path, + precond_post_multiply=fn_kind == "f_segment", ) launch_distributed_run( "apply_eigfn_to_query", @@ -178,6 +194,7 @@ def walk_query_phase1( method: str, lr_times_steps_per_segment: list[float], distributed: DistributedConfig, + preconditioner_paths: list[str] | None = None, ) -> list[Path]: """Phase 1: build query_grad_0, ..., query_grad_{L-1} by walking F_backward. @@ -203,6 +220,7 @@ def walk_query_phase1( lr_times_steps=lr_times_steps_per_segment[k], fn_kind="f_backward", distributed=distributed, + precond_path=preconditioner_paths[k] if preconditioner_paths else "", ) query_grad_paths[k - 1] = dst @@ -215,6 +233,7 @@ def walk_query_phase2( lr_times_steps_per_segment: list[float], query_grad_paths: list[Path], distributed: DistributedConfig, + preconditioner_paths: list[str] | None = None, ) -> list[Path]: """Phase 2: build query_grad_segment_0, ..., query_grad_segment_{L-1} via F_segment. @@ -238,6 +257,7 @@ def walk_query_phase2( lr_times_steps=lr_times_steps_per_segment[l], fn_kind="f_segment", distributed=distributed, + precond_path=preconditioner_paths[l] if preconditioner_paths else "", ) query_grad_segment_paths.append(dst) diff --git a/bergson/approx_unrolling/pipeline.py b/bergson/approx_unrolling/pipeline.py index a2853c8b..684b3891 100644 --- a/bergson/approx_unrolling/pipeline.py +++ b/bergson/approx_unrolling/pipeline.py @@ -37,7 +37,12 @@ PreprocessConfig, ) from ..config.config_io import save_run_config +from ..distributed import parent_barrier from ..utils.logger import get_logger +from .adam_preconditioner import ( + OPTIMIZER_STATE_FILE, + build_segment_preconditioners, +) from .approx_unrolling_math import ( compute_lr_times_steps_per_segment, score_per_segment_and_aggregate, @@ -104,6 +109,16 @@ def approx_unrolling_pipeline( assert hessian_cfg.ev_correction, "Approximate unrolling pipeline currently only " "supports EV correction on." + if approx_unrolling_cfg.use_adam_preconditioner: + for ckpt in approx_unrolling_cfg.checkpoints: + state_path = Path(ckpt) / OPTIMIZER_STATE_FILE + if not state_path.exists(): + raise FileNotFoundError( + f"use_adam_preconditioner requires {state_path}; write it " + "during training with TrainingConfig.save_optimizer_state " + "and copy it into the exported checkpoint dir." + ) + lr_times_steps_per_segment = compute_lr_times_steps_per_segment( approx_unrolling_cfg ) @@ -197,6 +212,24 @@ def approx_unrolling_pipeline( ) build(query_cfg, query_preprocess_cfg) + # ── Per-segment Adam preconditioners from the checkpoints' second moments + precond_paths: list[Path] = [] + if approx_unrolling_cfg.use_adam_preconditioner: + logger.info("Building per-segment Adam preconditioners...") + if index_cfg.distributed._node_rank == 0: + precond_paths = build_segment_preconditioners( + run_path=index_cfg.run_path, + method=hessian_cfg.method, + checkpoints=[str(c) for c in approx_unrolling_cfg.checkpoints], + segments=n_segments, + ) + else: + precond_paths = [ + Path(index_cfg.run_path) / f"segment_{l}" / "preconditioner.safetensors" + for l in range(n_segments) + ] + parent_barrier(index_cfg.distributed) + # ── Step 6: Phase 1 -- walk query backwards to get segment queries logger.info( f"Step 6/{_N_TOTAL_STEPS}: " @@ -204,11 +237,21 @@ def approx_unrolling_pipeline( f"query_grad_0..query_grad_(L-1)..." ) logger.info(f" lr_times_steps per segment: {lr_times_steps_per_segment}") + logger.info( + f" optimizer variant : " + f"{'preconditioned (Adam/AdamW)' if precond_paths else 'sgd'}" + + ( + f", momentum={approx_unrolling_cfg.momentum}" + if approx_unrolling_cfg.momentum + else "" + ) + ) query_grad_paths = walk_query_phase1( run_path=index_cfg.run_path, method=hessian_cfg.method, lr_times_steps_per_segment=lr_times_steps_per_segment, distributed=index_cfg.distributed, + preconditioner_paths=[str(p) for p in precond_paths] or None, ) # ── Step 7: Phase 2 -- Get per-ckpt queries from segment queries @@ -223,6 +266,7 @@ def approx_unrolling_pipeline( lr_times_steps_per_segment=lr_times_steps_per_segment, query_grad_paths=query_grad_paths, distributed=index_cfg.distributed, + preconditioner_paths=[str(p) for p in precond_paths] or None, ) # ── Step 8: Phase 3 -- per-segment scoring + sum diff --git a/bergson/config/config.py b/bergson/config/config.py index b11a117e..8d4cefa6 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -428,11 +428,6 @@ class ValidationConfig(TrainingConfig, ABC): permutation. These scores may be produced by items with fewer than 2 tokens.""" - save_retrained_models: bool = False - """When True, save each leave-k-out retrained model (HF format, weights + - tokenizer) to ``/retrained/subset_/`` so it can be reused for - later attribution queries without retraining. ~0.5 GB per subset for GPT-2.""" - subset_fraction: float = 0.0 """When > 0, each of the ``num_subsets`` leave-k-out subsets is an independent draw (without replacement within a subset, overlapping across @@ -796,6 +791,12 @@ class ApproxUnrollingConfig(Serializable): 2024, App. D.2). When ``None`` it's derived from the run if present or set to 0.0.""" + use_adam_preconditioner: bool = False + """Evaluate the unrolling eigenfunctions on a diagonal approximation of + the preconditioned Hessian P^1/2 H P^1/2 (the Adam/AdamW SOURCE variant, + Bae et al. 2024, Appendix C). Requires an ``optimizer.pt`` in every + checkpoint dir; see :mod:`bergson.approx_unrolling.adam_preconditioner`.""" + query: DataConfig = field(default_factory=DataConfig) """Query dataset spec; gradients computed at the final checkpoint.""" diff --git a/bergson/hessians/apply_hessian.py b/bergson/hessians/apply_hessian.py index 1f6f95a9..1e7674f4 100644 --- a/bergson/hessians/apply_hessian.py +++ b/bergson/hessians/apply_hessian.py @@ -14,7 +14,10 @@ from bergson.config import InversionConfig from bergson.data import column_offsets, create_index, load_gradients from bergson.distributed import init_dist -from bergson.hessians.preconditioner import FactoredPreconditioner +from bergson.hessians.preconditioner import ( + DiagonalFactoredPreconditioner, + FactoredPreconditioner, +) from bergson.utils.logger import get_logger from bergson.utils.utils import get_device, numpy_to_tensor @@ -34,9 +37,19 @@ class EkfacConfig: """When set, compress each module's IVHP output to a ``[p, p]`` Kronecker random projection (``P_S @ (H^-1 G) @ P_A^T``).""" projection_type: Literal["normal", "rademacher"] = "rademacher" - projection_scale: Literal["jl", "row_norm"] = "jl" """Must match the index being scored. See ``IndexConfig``.""" + precond_path: str = "" + """Safetensors of a diagonal optimizer preconditioner (module name -> + [out, in] grid). When set, ``apply_fn`` is evaluated elementwise on a + diagonal approximation of P^1/2 H P^1/2 in parameter space (the + Adam/AdamW approximate-unrolling variant) instead of on the eigenvalues + in the EKFAC eigenbasis. Requires ``apply_fn``.""" + precond_post_multiply: bool = False + """With ``precond_path``: multiply the applied function's output by the + preconditioner grid — the P^1/2 F(M) P^1/2 sandwich of the segment + eigenfunction. Leave False for the backward eigenfunction, whose + P^1/2 exp(-t M) P^-1/2 sandwich cancels.""" debug: bool = False @@ -84,14 +97,29 @@ def __init__( self.device = get_device(self.rank) def compute_ivhp_sharded(self): - preconditioner = FactoredPreconditioner.from_shards( - self.path, - rank=self.rank, - device=self.device, - inversion_cfg=None if self.apply_fn is not None else self.inversion_cfg, - apply_fn=self.apply_fn, - ev_correction=self.cfg.ev_correction, - ) + if self.cfg.precond_path: + if self.apply_fn is None: + raise ValueError("precond_path requires apply_fn.") + preconditioner = DiagonalFactoredPreconditioner.from_shards( + self.path, + self.cfg.precond_path, + rank=self.rank, + device=self.device, + apply_fn=self.apply_fn, + multiply_by_precond=self.cfg.precond_post_multiply, + ev_correction=self.cfg.ev_correction, + ) + else: + preconditioner = FactoredPreconditioner.from_shards( + self.path, + rank=self.rank, + device=self.device, + inversion_cfg=( + None if self.apply_fn is not None else self.inversion_cfg + ), + apply_fn=self.apply_fn, + ev_correction=self.cfg.ev_correction, + ) o_dims = { name: preconditioner.eigen_g[name].shape[1] diff --git a/bergson/hessians/preconditioner.py b/bergson/hessians/preconditioner.py index cffb619b..83524076 100644 --- a/bergson/hessians/preconditioner.py +++ b/bergson/hessians/preconditioner.py @@ -28,6 +28,7 @@ from typing import Protocol, runtime_checkable import torch +import torch.distributed as dist from safetensors.torch import load_file from torch import Tensor @@ -225,10 +226,10 @@ def _scale(self, name: str, g: Tensor) -> None: plus the globally-reduced means, then applied across shards in-place. """ lam = self.lambdas[name] + o, i = g.shape[1], lam.shape[1] if self.apply_fn is not None: inverse_eigvals = self.apply_fn(lam) else: - o, i = g.shape[1], lam.shape[1] mean = self.shard_computer.global_mean(lam, o * i) inversion = self.inversion_cfg.inversion if inversion == "factored_tikhonov": @@ -294,6 +295,178 @@ def apply(self, grads: dict[str, Tensor]) -> dict[str, Tensor]: return out +class DiagonalFactoredPreconditioner: + """Parameter-space diagonal eigenvalue-function preconditioner for the + preconditioned-optimizer (Adam/AdamW) approximate-unrolling variant + (Bae et al., 2024, Appendix C). + + With a diagonal optimizer preconditioner ``P`` the unrolling + eigenfunctions act on ``M = P^1/2 H P^1/2``, whose matrix exponential is + intractable in the K-FAC eigenbasis, so it is evaluated with a diagonal + Hessian approximation (Bae et al., Appendix D.2): no eigenbasis rotation, + just an elementwise multiplier ``apply_fn(p * d)`` (optionally times ``p``, + see ``multiply_by_precond``) where ``d = diag(H)`` is recovered from the + EKFAC factors as ``(Q_G ∘ Q_G) Λ (Q_A ∘ Q_A)^T``. + + ``apply_fn(sigma, mean)`` has the same contract as in + :class:`FactoredPreconditioner`: this rank's ``[c, I]`` row-shard of the + (here diagonal, parameter-space) eigenvalue grid plus its globally-reduced + mean. ``multiply_by_precond`` multiplies the result by ``p`` — the + ``P^1/2 F(M) P^1/2`` sandwich of the segment eigenfunction; the backward + eigenfunction's ``P^1/2 exp(-t M) P^-1/2`` sandwich cancels instead. + """ + + def __init__( + self, + eigen_a: dict[str, Tensor], + eigen_g: dict[str, Tensor], + lambdas: dict[str, Tensor], + precond: dict[str, Tensor], + *, + apply_fn, + multiply_by_precond: bool = False, + ): + self.eigen_a = eigen_a + self.eigen_g = eigen_g + self.lambdas = lambdas + self.precond = precond + self.apply_fn = apply_fn + self.multiply_by_precond = multiply_by_precond + self.shard_computer = ShardedMul() + self.logger = get_logger("DiagonalFactoredPreconditioner") + self._multipliers: dict[str, Tensor] = {} + + @classmethod + def from_shards( + cls, + hessian_path: str | Path, + precond_path: str | Path, + *, + rank: int, + device: str | torch.device, + apply_fn, + multiply_by_precond: bool = False, + ev_correction: bool = False, + ) -> "DiagonalFactoredPreconditioner": + """Load this rank's factor shards plus its row-shard of the full + parameter-space preconditioner grids saved at ``precond_path``.""" + lambda_dir = ( + "eigenvalue_correction_sharded" if ev_correction else "eigenvalue_sharded" + ) + eigen_a = _load_shard(hessian_path, "eigen_activation_sharded", rank, device) + eigen_g = _load_shard(hessian_path, "eigen_gradient_sharded", rank, device) + lambdas = _load_shard(hessian_path, lambda_dir, rank, device) + + full_precond = load_file(str(precond_path), device=str(device)) + sharder = ShardedMul() + precond = {} + for name, q_g in eigen_g.items(): + if name not in full_precond: + raise KeyError( + f"Module {name!r} has EKFAC factors but no preconditioner " + f"grid in {precond_path}; available: " + f"{sorted(full_precond.keys())}" + ) + p = full_precond[name].to(torch.float32) + o, i = q_g.shape[1], eigen_a[name].shape[1] + if p.shape != (o, i): + raise ValueError( + f"Preconditioner grid for {name!r} has shape " + f"{tuple(p.shape)}, expected [out, in] = ({o}, {i})." + ) + start, end = sharder.shard_bounds(o) + precond[name] = p[start:end] + return cls( + eigen_a, + eigen_g, + lambdas, + precond, + apply_fn=apply_fn, + multiply_by_precond=multiply_by_precond, + ) + + def _diag_hessian_shard(self, name: str) -> Tensor: + """This rank's ``[c_o, I]`` row-shard (rows = its block of the + parameter out-dim) of ``diag(H) = (Q_G ∘ Q_G) Λ (Q_A ∘ Q_A)^T``. + + The factors are row-sharded: ``Q_A [c_i, I']`` / ``Q_G [c_o, O']`` on + their parameter dims, ``Λ [c_o', I']`` on the eigen out-dim. Both + contractions run over a sharded dim, so each is a broadcast loop in + the style of :class:`ShardedMul`. + """ + q_a = self.eigen_a[name] # [c_i, I'] param rows, eigen cols + q_g = self.eigen_g[name] # [c_o, O'] param rows, eigen cols + lam = self.lambdas[name] # [c_o', I'] eigen rows, eigen cols + sc = self.shard_computer + n_eig_i = q_a.shape[1] + n_eig_o = q_g.shape[1] + + if not sc.dist: + # Single process: shards are the full factors. + return (q_g**2) @ lam @ (q_a**2).T + + # x_shard[o', i] = sum_i' lam[o', i'] * q_a[i, i']^2 for this rank's + # eigen-o' rows and ALL param-i columns (q_a's param rows are sharded). + x_shard = torch.empty(lam.shape[0], n_eig_i, device=lam.device, dtype=lam.dtype) + for rank_index in range(sc.world_size): + start, end = sc.shard_bounds(n_eig_i, rank_index) + if rank_index == sc.rank: + shard = q_a + else: + shard = torch.empty( + (end - start, q_a.shape[1]), device=q_a.device, dtype=q_a.dtype + ) + dist.broadcast(shard, src=rank_index) + x_shard[:, start:end] = lam @ (shard**2).T + if sc.rank != rank_index: + del shard + + # d_shard[o, i] = sum_o' q_g[o, o']^2 * x[o', i]; x rows (eigen-o') are + # sharded across ranks, q_g's eigen-o' columns are full. + d_shard = torch.zeros(q_g.shape[0], n_eig_i, device=q_g.device, dtype=q_g.dtype) + for rank_index in range(sc.world_size): + start, end = sc.shard_bounds(n_eig_o, rank_index) + if rank_index == sc.rank: + shard = x_shard + else: + shard = torch.empty( + (end - start, n_eig_i), device=x_shard.device, dtype=x_shard.dtype + ) + dist.broadcast(shard, src=rank_index) + d_shard += (q_g[:, start:end] ** 2) @ shard + if sc.rank != rank_index: + del shard + return d_shard + + def _multiplier(self, name: str) -> Tensor: + """This rank's ``[c_o, I]`` row-shard of the elementwise multiplier, + cached after the first batch (it is gradient-independent).""" + if name not in self._multipliers: + p = self.precond[name] + sigma = p * self._diag_hessian_shard(name) + multiplier = self.apply_fn(sigma) + if self.multiply_by_precond: + multiplier = multiplier * p + self._multipliers[name] = multiplier + return self._multipliers[name] + + def apply(self, grads: dict[str, Tensor]) -> dict[str, Tensor]: + """Return ``grads`` scaled elementwise in parameter space; modules + absent from the factors pass through unchanged.""" + out: dict[str, Tensor] = {} + for name, flat in grads.items(): + if name not in self.eigen_a: + out[name] = flat + continue + o = self.eigen_g[name].shape[1] + i = self.eigen_a[name].shape[1] + g = flat.to(self.lambdas[name].device, torch.float32).view(-1, o, i) + self.shard_computer.scale_rows_in_place(g, self._multiplier(name)) + self.logger.debug("%s: scaled elementwise by f(p * diag(H))", name) + out[name] = g.reshape(flat.shape[0], -1).to(flat.dtype) + return out + + def is_factored_hessian(hessian_path: str | Path) -> bool: """Whether ``hessian_path`` holds a Kronecker-factored (EKFAC) Hessian rather than a dense one — detected by the presence of the eigenvector shard directory. diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index 9f0e5979..aff5d580 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -263,9 +263,11 @@ def worker( else None ), ) - if getattr(run_cfg, "save_optimizer_state", "none") != "none" and global_rank == 0: + # Called on every rank: FSDP moments are DTensors whose gather is a + # collective; rank 0 writes inside. + if getattr(run_cfg, "save_optimizer_state", "none") != "none": save_second_moments_as_optimizer_pt( - model, # type: ignore[reportArgumentType] + model, fwd_state.opt_state, os.path.join(run_cfg.run_path, "optimizer.pt"), ) diff --git a/tests/test_distributed_magic.py b/tests/test_distributed_magic.py index 06492026..d39e23cb 100644 --- a/tests/test_distributed_magic.py +++ b/tests/test_distributed_magic.py @@ -239,3 +239,23 @@ def run(name: str, *, fsdp: bool, grad_accum: int) -> torch.Tensor: f"FSDP and DDP scores differ under dropout: {fsdp_ddp_diff:.2e} " f"(scale {scale:.2e}) — the ranks are likely drawing different masks" ) + + +@requires_multi_gpu +def test_save_optimizer_state_completes_under_fsdp(tmp_path): + """save_optimizer_state must not hang under FSDP. + + Gathering each rank's sharded second moments (DTensor.full_tensor()) is a + collective, so every rank must call save_second_moments_as_optimizer_pt; + if worker() ever gates that call behind `global_rank == 0` again, this + hangs instead of failing cleanly, since the rank-0 all-gather then waits + forever for peers that never join it. + """ + cfg = magic_cfg(f"{tmp_path}/fsdp_opt", fsdp=True, clip=False) + cfg.save_optimizer_state = True + run_magic(cfg) + + optimizer_pt = torch.load(f"{tmp_path}/fsdp_opt/optimizer.pt", weights_only=False) + assert optimizer_pt["state"], "optimizer.pt has no second-moment entries" + for entry in optimizer_pt["state"].values(): + assert torch.isfinite(entry["exp_avg_sq"]).all() diff --git a/tests/test_source_optimizer_variants.py b/tests/test_source_optimizer_variants.py new file mode 100644 index 00000000..b7a3a80e --- /dev/null +++ b/tests/test_source_optimizer_variants.py @@ -0,0 +1,312 @@ +"""Tests for the SOURCE (approx-unrolling) optimizer variants: SGD heavy-ball +momentum lr scaling and the Adam/AdamW diagonal-preconditioner eigenfunction +path (Bae et al., 2024, Appendix C / D.2). +""" + +import pytest +import torch +from safetensors.torch import save_file + +from bergson.approx_unrolling.approx_unrolling_math import ( + compute_lr_times_steps_per_segment, + f_backward, + f_segment, +) +from bergson.config import ApproxUnrollingConfig +from bergson.hessians.preconditioner import DiagonalFactoredPreconditioner + +MODULE = "lm_head" +OUT_DIM, IN_DIM = 6, 4 +LR_TIMES_STEPS = 0.05 + + +def test_momentum_scales_lr_times_steps(): + """SGDm terminal velocity: lr*K scaled by 1/(1-beta).""" + base_cfg = ApproxUnrollingConfig( + checkpoints=["a", "b"], + segments=2, + lr_list=[1e-3, 2e-3], + step_size_list=[10, 20], + ) + momentum_cfg = ApproxUnrollingConfig( + checkpoints=["a", "b"], + segments=2, + lr_list=[1e-3, 2e-3], + step_size_list=[10, 20], + momentum=0.9, + ) + base = compute_lr_times_steps_per_segment(base_cfg) + scaled = compute_lr_times_steps_per_segment(momentum_cfg) + assert scaled == pytest.approx([10 * b for b in base]) + + +def test_momentum_out_of_range_raises(): + for bad in (1.0, -0.1): + bad_cfg = ApproxUnrollingConfig( + checkpoints=["a"], + segments=1, + lr_list=[1e-3], + step_size_list=[10], + momentum=bad, + ) + with pytest.raises(ValueError, match="momentum"): + compute_lr_times_steps_per_segment(bad_cfg) + + +def _random_factors(): + """Random EKFAC-style factors: orthogonal eigenvectors, non-negative + eigenvalue grid.""" + torch.manual_seed(0) + q_a, _ = torch.linalg.qr(torch.randn(IN_DIM, IN_DIM, dtype=torch.float64)) + q_g, _ = torch.linalg.qr(torch.randn(OUT_DIM, OUT_DIM, dtype=torch.float64)) + lam = torch.rand(OUT_DIM, IN_DIM, dtype=torch.float64) + return q_a.float().contiguous(), q_g.float().contiguous(), lam.float() + + +def _write_factor_shards(tmp_path, q_a, q_g, lam, precond): + """Write the single-shard on-disk layout DiagonalFactoredPreconditioner + loads from, plus the preconditioner grid.""" + for sub, tensor in [ + ("eigen_activation_sharded", q_a), + ("eigen_gradient_sharded", q_g), + ("eigenvalue_correction_sharded", lam), + ]: + (tmp_path / sub).mkdir() + save_file({MODULE: tensor}, str(tmp_path / sub / "shard_0.safetensors")) + precond_path = tmp_path / "precond.safetensors" + save_file({MODULE: precond}, str(precond_path)) + return precond_path + + +def _reference_diag_hessian(q_a, q_g, lam): + """diag(H) via the dense Kronecker Hessian — independent of the + implementation's factored formula. + + The [OUT, IN] grid flattens row-major to vec index o*IN + i, matching + kron(q_g, q_a)'s row ordering, so H = kron(Q_G, Q_A) diag(vec Λ) + kron(Q_G, Q_A)^T. + """ + q_kron = torch.kron(q_g.double(), q_a.double()) + h = q_kron @ torch.diag(lam.double().flatten()) @ q_kron.T + return torch.diagonal(h).reshape(OUT_DIM, IN_DIM).float() + + +@pytest.mark.parametrize("fn_kind", ["f_backward", "f_segment"]) +def test_diagonal_preconditioner_matches_dense_reference(tmp_path, fn_kind): + """The diagonal path multiplies gradients elementwise by + f(p * diag(H)) — times p for f_segment's P^1/2 F(M) P^1/2 sandwich — + with diag(H) recovered exactly from the EKFAC factors.""" + q_a, q_g, lam = _random_factors() + precond = torch.rand(OUT_DIM, IN_DIM) + 0.5 + precond_path = _write_factor_shards(tmp_path, q_a, q_g, lam, precond) + + fn = {"f_backward": f_backward, "f_segment": f_segment}[fn_kind](LR_TIMES_STEPS) + preconditioner = DiagonalFactoredPreconditioner.from_shards( + tmp_path, + precond_path, + rank=0, + device="cpu", + apply_fn=fn, + multiply_by_precond=fn_kind == "f_segment", + ev_correction=True, + ) + + grads = {MODULE: torch.randn(3, OUT_DIM * IN_DIM)} + out = preconditioner.apply({k: v.clone() for k, v in grads.items()})[MODULE] + + sigma = precond * _reference_diag_hessian(q_a, q_g, lam) + if fn_kind == "f_backward": + multiplier = torch.exp(-LR_TIMES_STEPS * sigma) + else: + multiplier = precond * (-torch.expm1(-LR_TIMES_STEPS * sigma)) / sigma + expected = grads[MODULE].view(3, OUT_DIM, IN_DIM) * multiplier + torch.testing.assert_close( + out.view(3, OUT_DIM, IN_DIM), expected, rtol=1e-4, atol=1e-6 + ) + + +def test_f_segment_zero_eigenvalue_limit(tmp_path): + """With Λ = 0 the segment multiplier hits its lr*K limit, times the + preconditioner: lr*K * p, with no NaN/inf from the 0/0.""" + q_a, q_g, _ = _random_factors() + lam = torch.zeros(OUT_DIM, IN_DIM) + precond = torch.rand(OUT_DIM, IN_DIM) + 0.5 + precond_path = _write_factor_shards(tmp_path, q_a, q_g, lam, precond) + + preconditioner = DiagonalFactoredPreconditioner.from_shards( + tmp_path, + precond_path, + rank=0, + device="cpu", + apply_fn=f_segment(LR_TIMES_STEPS), + multiply_by_precond=True, + ev_correction=True, + ) + grads = {MODULE: torch.ones(1, OUT_DIM * IN_DIM)} + out = preconditioner.apply(grads)[MODULE].view(OUT_DIM, IN_DIM) + torch.testing.assert_close(out, LR_TIMES_STEPS * precond) + + +def test_precond_shape_mismatch_raises(tmp_path): + q_a, q_g, lam = _random_factors() + bad_precond = torch.rand(OUT_DIM + 1, IN_DIM) + precond_path = _write_factor_shards(tmp_path, q_a, q_g, lam, bad_precond) + with pytest.raises(ValueError, match="shape"): + DiagonalFactoredPreconditioner.from_shards( + tmp_path, + precond_path, + rank=0, + device="cpu", + apply_fn=f_backward(LR_TIMES_STEPS), + ev_correction=True, + ) + + +def test_precond_missing_module_raises(tmp_path): + q_a, q_g, lam = _random_factors() + for sub, tensor in [ + ("eigen_activation_sharded", q_a), + ("eigen_gradient_sharded", q_g), + ("eigenvalue_correction_sharded", lam), + ]: + (tmp_path / sub).mkdir() + save_file({MODULE: tensor}, str(tmp_path / sub / "shard_0.safetensors")) + precond_path = tmp_path / "precond.safetensors" + save_file({"other_module": torch.rand(OUT_DIM, IN_DIM)}, str(precond_path)) + with pytest.raises(KeyError, match=MODULE): + DiagonalFactoredPreconditioner.from_shards( + tmp_path, + precond_path, + rank=0, + device="cpu", + apply_fn=f_backward(LR_TIMES_STEPS), + ev_correction=True, + ) + + +def test_build_segment_preconditioners(tmp_path): + """Per-segment P built from the checkpoints' optimizer.pt files: + bias-corrected via the stored step/betas, index-mapped to param names, + suffix-matched to the factor modules, oriented to [out, in] (transposed + storage flipped), averaged within the segment, and eps-transformed.""" + from transformers import AutoModelForCausalLM, GPT2Config + + from bergson.approx_unrolling.adam_preconditioner import ( + build_segment_preconditioners, + ) + + torch.manual_seed(0) + tiny_cfg = GPT2Config(n_layer=1, n_embd=4, n_head=2, n_positions=8, vocab_size=16) + module, out_dim, in_dim = "h.0.attn.c_attn", 12, 4 + + run = tmp_path / "run" + seg_kfac = run / "segment_0" / "kfac" + seg_kfac.mkdir(parents=True) + for sub, cols in [ + ("eigen_activation_sharded", in_dim), + ("eigen_gradient_sharded", out_dim), + ("eigenvalue_correction_sharded", in_dim), + ]: + (seg_kfac / sub).mkdir() + save_file( + {module: torch.rand(2, cols)}, str(seg_kfac / sub / "shard_0.safetensors") + ) + + tiny_model = AutoModelForCausalLM.from_config(tiny_cfg) + param_names = [n for n, _ in tiny_model.named_parameters()] + idx = param_names.index(f"transformer.{module}.weight") + + # Two checkpoints with raw (uncorrected) moments stored transposed + # ([in, out], like GPT-2 Conv1D params), at different steps. + beta2, eps_root, eps = 0.975, 1e-6, 1e-8 + steps, nus = [5, 9], [torch.rand(in_dim, out_dim) for _ in range(2)] + ckpts = [] + for step, nu in zip(steps, nus): + ckpt = tmp_path / f"models_step_{step}" + ckpt.mkdir() + tiny_cfg.save_pretrained(ckpt) + # One checkpoint keyed positionally (legacy fallback), the other under + # a bogus index with param_name recorded (the FSDP-scrambled case). + if step == steps[0]: + entry = {idx: {"exp_avg_sq": nu, "step": torch.tensor(step)}} + else: + entry = { + 999: { + "exp_avg_sq": nu, + "step": torch.tensor(step), + "param_name": f"transformer.{module}.weight", + } + } + torch.save( + { + "state": entry, + "param_groups": [ + { + "params": list(entry), + "betas": (0.9, beta2), + "eps": eps, + "eps_root": eps_root, + } + ], + }, + ckpt / "optimizer.pt", + ) + ckpts.append(str(ckpt)) + + paths = build_segment_preconditioners( + run_path=run, + method="kfac", + checkpoints=ckpts, + segments=1, + ) + assert paths == [run / "segment_0" / "preconditioner.safetensors"] + from safetensors.torch import load_file + + precond = load_file(str(paths[0]))[module] + v_hats = [nu / (1 - beta2**step) for step, nu in zip(steps, nus)] + v_bar = (v_hats[0] + v_hats[1]).T / 2 + expected = 1.0 / ((v_bar + eps_root).sqrt() + eps) + assert precond.shape == (out_dim, in_dim) + torch.testing.assert_close(precond, expected) + + +def test_optimizer_pt_snapshot_fields(tmp_path): + """save_second_moments_as_optimizer_pt stores the standard step and betas + fields when given, making snapshot exports self-describing for the SOURCE + Adam variant's bias correction.""" + import torch.nn as nn + import torchopt + + from bergson.utils.load_from_optimizer import save_second_moments_as_optimizer_pt + + class TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.blk = nn.Linear(3, 2, bias=False) + self.head = nn.Linear(2, 4, bias=False) + + model = TinyModel() + params = dict(model.named_parameters(remove_duplicate=False)) + opt_state = torchopt.adamw(1e-3, betas=(0.9, 0.975)).init(params) + adam_state = next(s for s in opt_state if hasattr(s, "nu")) + for nu, count in zip(adam_state.nu, adam_state.count): + nu.copy_(torch.rand_like(nu)) + count.fill_(7) + + path = tmp_path / "step_7.optimizer.pt" + n = save_second_moments_as_optimizer_pt( + model, opt_state, path, step=7, betas=(0.9, 0.975), eps=1e-8, eps_root=1e-6 + ) + assert n == 2 + optimizer_pt = torch.load(path, weights_only=False) + assert optimizer_pt["param_groups"][0]["betas"] == (0.9, 0.975) + assert optimizer_pt["param_groups"][0]["eps"] == 1e-8 + assert optimizer_pt["param_groups"][0]["eps_root"] == 1e-6 + names = [n for n, _ in model.named_parameters()] + # nu lists are in sorted(params) order; blk.weight sorts before head.weight. + for idx, entry in optimizer_pt["state"].items(): + assert int(entry["step"].item()) == 7 + # param_name recorded per entry: FSDP-scrambled indices stay readable. + assert entry["param_name"] == names[idx] + nu_idx = sorted(params).index(names[idx]) + torch.testing.assert_close(entry["exp_avg_sq"], adam_state.nu[nu_idx]) From 85f434fdd332daa0879cdba8a34fb2e61a58c079 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Sat, 25 Jul 2026 09:49:39 +0000 Subject: [PATCH 02/10] refactor: clarify ADAM SOURCE naming; cover the applicator wiring Spell out the abbreviated preconditioner names so they read alongside hessian_path: EkfacConfig.precond_path/precond_post_multiply -> preconditioner_path/preconditioner_post_multiply, and DiagonalFactoredPreconditioner's multiply_by_precond -> multiply_by_preconditioner. Rename Trainer.train's snapshot_optimizer_hparams -> optimizer_hparams. Drop the getattr() guards on save_optimizer_state/save_retrained_models in magic's worker: both are TrainingConfig fields now, and worker takes a TrainingConfig. Trim the DiagonalFactoredPreconditioner and adam_preconditioner docs to the formula and the citation; the "called on every rank" FSDP comment is already pinned by test_save_optimizer_state_completes_under_fsdp, whose docstring explains the hang it prevents. The only untested part of this feature was the config wiring, so drive a real EkfacApplicator with preconditioner_path set through an on-disk gradient index and compare against DiagonalFactoredPreconditioner.apply for both eigenfunctions. That is what catches a rename that fails to reach compute_ivhp_sharded. CPU-only, 0.25s. --- .../approx_unrolling/adam_preconditioner.py | 22 ++--- .../approx_unrolling/approx_unrolling_math.py | 12 +-- bergson/approx_unrolling/pipeline.py | 12 +-- bergson/hessians/apply_hessian.py | 16 ++-- bergson/hessians/preconditioner.py | 62 ++++++------- bergson/magic/cli.py | 14 ++- tests/test_source_optimizer_variants.py | 92 +++++++++++++++---- 7 files changed, 144 insertions(+), 86 deletions(-) diff --git a/bergson/approx_unrolling/adam_preconditioner.py b/bergson/approx_unrolling/adam_preconditioner.py index 30ca2187..60edb276 100644 --- a/bergson/approx_unrolling/adam_preconditioner.py +++ b/bergson/approx_unrolling/adam_preconditioner.py @@ -1,19 +1,15 @@ """Per-segment Adam/AdamW diagonal preconditioners for the preconditioned SOURCE (approx-unrolling) variant. -Loads optimizer.pt from the checkpoint dirs. You can use -``TrainingConfig.save_optimizer_state`` to save them with the Bergson -Trainer. They contain the second moment buffer, ``step``, and ``betas``). -:func:`build_segment_preconditioners` bias-corrects the moments -(v_hat = exp_avg_sq / (1 - beta2^step)), matches them to the segment's -EKFAC factor modules, orients them to bergson's ``[out, in]`` gradient -grids (params stored transposed, e.g. GPT-2 Conv1D, are flipped), averages -within each segment (the same stationary-within-segment averaging used for -the K-FAC factors), and writes the diagonal preconditioner +:func:`build_segment_preconditioners` reads each checkpoint dir's +``optimizer.pt`` (written by ``TrainingConfig.save_optimizer_state``), +bias-corrects the second moments, averages them within each segment as the +K-FAC factors are, and writes P = 1 / (sqrt(v_hat_bar + eps_root) + eps) -to ``/segment_/preconditioner.safetensors``. +to ``/segment_/preconditioner.safetensors``, oriented to bergson's +``[out, in]`` gradient grids. """ from glob import glob @@ -175,7 +171,7 @@ def build_segment_preconditioners( grids = _module_grids(base / f"segment_{seg}" / method) matches = _match_params(grids, list(v_hat_sum)) - precond = {} + preconditioner = {} for module, param_name in matches.items(): try: layer = reference_model.get_submodule( @@ -185,9 +181,9 @@ def build_segment_preconditioners( layer = None v_hat_bar = v_hat_sum[param_name] / len(seg_ckpts) p = 1.0 / ((v_hat_bar + eps_root).sqrt() + eps) - precond[module] = _orient(p, grids[module], module, layer) + preconditioner[module] = _orient(p, grids[module], module, layer) out_path = base / f"segment_{seg}" / "preconditioner.safetensors" - save_file(precond, str(out_path)) + save_file(preconditioner, str(out_path)) out_paths.append(out_path) return out_paths diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index a96e9e30..b607392d 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -146,13 +146,13 @@ def apply_eigfn_to_query( lr_times_steps: float, fn_kind: str, distributed: DistributedConfig, - precond_path: str = "", + preconditioner_path: str = "", ) -> None: """Apply F_segment or F_backward of one segment to a stored query gradient. ``fn_kind`` is "f_segment" or "f_backward". The segment eigenvalues are already checkpoint-averaged (expected eigenvalues), so the eigenfunction is - applied to them directly. ``precond_path`` selects the + applied to them directly. ``preconditioner_path`` selects the preconditioned-optimizer variant: the eigenfunction is evaluated on a diagonal approximation of P^1/2 H P^1/2 in parameter space, with F_segment's output additionally multiplied by P (its P^1/2 . P^1/2 @@ -162,8 +162,8 @@ def apply_eigfn_to_query( gradient_path=str(src_grad_path), run_path=str(dst_grad_path), ev_correction=True, - precond_path=precond_path, - precond_post_multiply=fn_kind == "f_segment", + preconditioner_path=preconditioner_path, + preconditioner_post_multiply=fn_kind == "f_segment", ) launch_distributed_run( "apply_eigfn_to_query", @@ -220,7 +220,7 @@ def walk_query_phase1( lr_times_steps=lr_times_steps_per_segment[k], fn_kind="f_backward", distributed=distributed, - precond_path=preconditioner_paths[k] if preconditioner_paths else "", + preconditioner_path=preconditioner_paths[k] if preconditioner_paths else "", ) query_grad_paths[k - 1] = dst @@ -257,7 +257,7 @@ def walk_query_phase2( lr_times_steps=lr_times_steps_per_segment[l], fn_kind="f_segment", distributed=distributed, - precond_path=preconditioner_paths[l] if preconditioner_paths else "", + preconditioner_path=preconditioner_paths[l] if preconditioner_paths else "", ) query_grad_segment_paths.append(dst) diff --git a/bergson/approx_unrolling/pipeline.py b/bergson/approx_unrolling/pipeline.py index 684b3891..dbc9902d 100644 --- a/bergson/approx_unrolling/pipeline.py +++ b/bergson/approx_unrolling/pipeline.py @@ -213,18 +213,18 @@ def approx_unrolling_pipeline( build(query_cfg, query_preprocess_cfg) # ── Per-segment Adam preconditioners from the checkpoints' second moments - precond_paths: list[Path] = [] + preconditioner_paths: list[Path] = [] if approx_unrolling_cfg.use_adam_preconditioner: logger.info("Building per-segment Adam preconditioners...") if index_cfg.distributed._node_rank == 0: - precond_paths = build_segment_preconditioners( + preconditioner_paths = build_segment_preconditioners( run_path=index_cfg.run_path, method=hessian_cfg.method, checkpoints=[str(c) for c in approx_unrolling_cfg.checkpoints], segments=n_segments, ) else: - precond_paths = [ + preconditioner_paths = [ Path(index_cfg.run_path) / f"segment_{l}" / "preconditioner.safetensors" for l in range(n_segments) ] @@ -239,7 +239,7 @@ def approx_unrolling_pipeline( logger.info(f" lr_times_steps per segment: {lr_times_steps_per_segment}") logger.info( f" optimizer variant : " - f"{'preconditioned (Adam/AdamW)' if precond_paths else 'sgd'}" + f"{'preconditioned (Adam/AdamW)' if preconditioner_paths else 'sgd'}" + ( f", momentum={approx_unrolling_cfg.momentum}" if approx_unrolling_cfg.momentum @@ -251,7 +251,7 @@ def approx_unrolling_pipeline( method=hessian_cfg.method, lr_times_steps_per_segment=lr_times_steps_per_segment, distributed=index_cfg.distributed, - preconditioner_paths=[str(p) for p in precond_paths] or None, + preconditioner_paths=[str(p) for p in preconditioner_paths] or None, ) # ── Step 7: Phase 2 -- Get per-ckpt queries from segment queries @@ -266,7 +266,7 @@ def approx_unrolling_pipeline( lr_times_steps_per_segment=lr_times_steps_per_segment, query_grad_paths=query_grad_paths, distributed=index_cfg.distributed, - preconditioner_paths=[str(p) for p in precond_paths] or None, + preconditioner_paths=[str(p) for p in preconditioner_paths] or None, ) # ── Step 8: Phase 3 -- per-segment scoring + sum diff --git a/bergson/hessians/apply_hessian.py b/bergson/hessians/apply_hessian.py index 1e7674f4..ed17fb28 100644 --- a/bergson/hessians/apply_hessian.py +++ b/bergson/hessians/apply_hessian.py @@ -39,15 +39,15 @@ class EkfacConfig: projection_type: Literal["normal", "rademacher"] = "rademacher" projection_scale: Literal["jl", "row_norm"] = "jl" """Must match the index being scored. See ``IndexConfig``.""" - precond_path: str = "" + preconditioner_path: str = "" """Safetensors of a diagonal optimizer preconditioner (module name -> [out, in] grid). When set, ``apply_fn`` is evaluated elementwise on a diagonal approximation of P^1/2 H P^1/2 in parameter space (the Adam/AdamW approximate-unrolling variant) instead of on the eigenvalues in the EKFAC eigenbasis. Requires ``apply_fn``.""" - precond_post_multiply: bool = False - """With ``precond_path``: multiply the applied function's output by the - preconditioner grid — the P^1/2 F(M) P^1/2 sandwich of the segment + preconditioner_post_multiply: bool = False + """With ``preconditioner_path``: multiply the applied function's output by + the preconditioner grid — the P^1/2 F(M) P^1/2 sandwich of the segment eigenfunction. Leave False for the backward eigenfunction, whose P^1/2 exp(-t M) P^-1/2 sandwich cancels.""" debug: bool = False @@ -97,16 +97,16 @@ def __init__( self.device = get_device(self.rank) def compute_ivhp_sharded(self): - if self.cfg.precond_path: + if self.cfg.preconditioner_path: if self.apply_fn is None: - raise ValueError("precond_path requires apply_fn.") + raise ValueError("preconditioner_path requires apply_fn.") preconditioner = DiagonalFactoredPreconditioner.from_shards( self.path, - self.cfg.precond_path, + self.cfg.preconditioner_path, rank=self.rank, device=self.device, apply_fn=self.apply_fn, - multiply_by_precond=self.cfg.precond_post_multiply, + multiply_by_preconditioner=self.cfg.preconditioner_post_multiply, ev_correction=self.cfg.ev_correction, ) else: diff --git a/bergson/hessians/preconditioner.py b/bergson/hessians/preconditioner.py index 83524076..ef81751e 100644 --- a/bergson/hessians/preconditioner.py +++ b/bergson/hessians/preconditioner.py @@ -296,24 +296,14 @@ def apply(self, grads: dict[str, Tensor]) -> dict[str, Tensor]: class DiagonalFactoredPreconditioner: - """Parameter-space diagonal eigenvalue-function preconditioner for the - preconditioned-optimizer (Adam/AdamW) approximate-unrolling variant - (Bae et al., 2024, Appendix C). - - With a diagonal optimizer preconditioner ``P`` the unrolling - eigenfunctions act on ``M = P^1/2 H P^1/2``, whose matrix exponential is - intractable in the K-FAC eigenbasis, so it is evaluated with a diagonal - Hessian approximation (Bae et al., Appendix D.2): no eigenbasis rotation, - just an elementwise multiplier ``apply_fn(p * d)`` (optionally times ``p``, - see ``multiply_by_precond``) where ``d = diag(H)`` is recovered from the - EKFAC factors as ``(Q_G ∘ Q_G) Λ (Q_A ∘ Q_A)^T``. - - ``apply_fn(sigma, mean)`` has the same contract as in - :class:`FactoredPreconditioner`: this rank's ``[c, I]`` row-shard of the - (here diagonal, parameter-space) eigenvalue grid plus its globally-reduced - mean. ``multiply_by_precond`` multiplies the result by ``p`` — the - ``P^1/2 F(M) P^1/2`` sandwich of the segment eigenfunction; the backward - eigenfunction's ``P^1/2 exp(-t M) P^-1/2`` sandwich cancels instead. + """Applies ``apply_fn`` elementwise in parameter space, on the diagonal of + ``P^1/2 H P^1/2`` for an optimizer preconditioner ``P`` — the Adam/AdamW + approximate-unrolling variant (Bae et al., 2024, Appendices C and D.2). + + There is no eigenbasis rotation: the multiplier is ``apply_fn(p * d)``, + with ``d = diag(H) = (Q_G ∘ Q_G) Λ (Q_A ∘ Q_A)^T`` from the EKFAC factors, + times ``p`` when ``multiply_by_preconditioner`` is set (F_segment's + ``P^1/2 F(M) P^1/2`` sandwich; F_backward's cancels). """ def __init__( @@ -321,17 +311,17 @@ def __init__( eigen_a: dict[str, Tensor], eigen_g: dict[str, Tensor], lambdas: dict[str, Tensor], - precond: dict[str, Tensor], + preconditioner: dict[str, Tensor], *, apply_fn, - multiply_by_precond: bool = False, + multiply_by_preconditioner: bool = False, ): self.eigen_a = eigen_a self.eigen_g = eigen_g self.lambdas = lambdas - self.precond = precond + self.preconditioner = preconditioner self.apply_fn = apply_fn - self.multiply_by_precond = multiply_by_precond + self.multiply_by_preconditioner = multiply_by_preconditioner self.shard_computer = ShardedMul() self.logger = get_logger("DiagonalFactoredPreconditioner") self._multipliers: dict[str, Tensor] = {} @@ -340,16 +330,16 @@ def __init__( def from_shards( cls, hessian_path: str | Path, - precond_path: str | Path, + preconditioner_path: str | Path, *, rank: int, device: str | torch.device, apply_fn, - multiply_by_precond: bool = False, + multiply_by_preconditioner: bool = False, ev_correction: bool = False, ) -> "DiagonalFactoredPreconditioner": """Load this rank's factor shards plus its row-shard of the full - parameter-space preconditioner grids saved at ``precond_path``.""" + parameter-space preconditioner grids saved at ``preconditioner_path``.""" lambda_dir = ( "eigenvalue_correction_sharded" if ev_correction else "eigenvalue_sharded" ) @@ -357,17 +347,17 @@ def from_shards( eigen_g = _load_shard(hessian_path, "eigen_gradient_sharded", rank, device) lambdas = _load_shard(hessian_path, lambda_dir, rank, device) - full_precond = load_file(str(precond_path), device=str(device)) + full_grids = load_file(str(preconditioner_path), device=str(device)) sharder = ShardedMul() - precond = {} + preconditioner = {} for name, q_g in eigen_g.items(): - if name not in full_precond: + if name not in full_grids: raise KeyError( f"Module {name!r} has EKFAC factors but no preconditioner " - f"grid in {precond_path}; available: " - f"{sorted(full_precond.keys())}" + f"grid in {preconditioner_path}; available: " + f"{sorted(full_grids.keys())}" ) - p = full_precond[name].to(torch.float32) + p = full_grids[name].to(torch.float32) o, i = q_g.shape[1], eigen_a[name].shape[1] if p.shape != (o, i): raise ValueError( @@ -375,14 +365,14 @@ def from_shards( f"{tuple(p.shape)}, expected [out, in] = ({o}, {i})." ) start, end = sharder.shard_bounds(o) - precond[name] = p[start:end] + preconditioner[name] = p[start:end] return cls( eigen_a, eigen_g, lambdas, - precond, + preconditioner, apply_fn=apply_fn, - multiply_by_precond=multiply_by_precond, + multiply_by_preconditioner=multiply_by_preconditioner, ) def _diag_hessian_shard(self, name: str) -> Tensor: @@ -442,10 +432,10 @@ def _multiplier(self, name: str) -> Tensor: """This rank's ``[c_o, I]`` row-shard of the elementwise multiplier, cached after the first batch (it is gradient-independent).""" if name not in self._multipliers: - p = self.precond[name] + p = self.preconditioner[name] sigma = p * self._diag_hessian_shard(name) multiplier = self.apply_fn(sigma) - if self.multiply_by_precond: + if self.multiply_by_preconditioner: multiplier = multiplier * p self._multipliers[name] = multiplier return self._multipliers[name] diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index aff5d580..b3de45c3 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -253,12 +253,17 @@ def worker( fsdp=run_cfg.fsdp, max_grad_norm=run_cfg.max_grad_norm, grad_accum_steps=run_cfg.grad_accum_steps, +<<<<<<< HEAD optimizer_cfg=( +======= + optimizer_hparams=( +>>>>>>> 5139d31a (refactor: clarify ADAM SOURCE naming; cover the applicator wiring) dict( betas=(run_cfg.adam_beta1, run_cfg.adam_beta2), eps=run_cfg.adam_eps, eps_root=run_cfg.eps_root, ) +<<<<<<< HEAD if getattr(run_cfg, "save_optimizer_state", "none") == "all" else None ), @@ -266,13 +271,20 @@ def worker( # Called on every rank: FSDP moments are DTensors whose gather is a # collective; rank 0 writes inside. if getattr(run_cfg, "save_optimizer_state", "none") != "none": +======= + if run_cfg.save_optimizer_state + else None + ), + ) + if run_cfg.save_optimizer_state: +>>>>>>> 5139d31a (refactor: clarify ADAM SOURCE naming; cover the applicator wiring) save_second_moments_as_optimizer_pt( model, fwd_state.opt_state, os.path.join(run_cfg.run_path, "optimizer.pt"), ) - if getattr(run_cfg, "save_retrained_models", False) and global_rank == 0: + if run_cfg.save_retrained_models and global_rank == 0: # Persist the fully-trained (no leave-out) model alongside the per-subset # models, so a later evaluate_retrained run can measure the query # baseline from the bank itself instead of a separate base_model_dir. diff --git a/tests/test_source_optimizer_variants.py b/tests/test_source_optimizer_variants.py index b7a3a80e..71866a23 100644 --- a/tests/test_source_optimizer_variants.py +++ b/tests/test_source_optimizer_variants.py @@ -3,6 +3,7 @@ path (Bae et al., 2024, Appendix C / D.2). """ +import numpy as np import pytest import torch from safetensors.torch import save_file @@ -13,6 +14,8 @@ f_segment, ) from bergson.config import ApproxUnrollingConfig +from bergson.data import create_index, load_module_gradients +from bergson.hessians.apply_hessian import EkfacApplicator, EkfacConfig from bergson.hessians.preconditioner import DiagonalFactoredPreconditioner MODULE = "lm_head" @@ -73,9 +76,9 @@ def _write_factor_shards(tmp_path, q_a, q_g, lam, precond): ]: (tmp_path / sub).mkdir() save_file({MODULE: tensor}, str(tmp_path / sub / "shard_0.safetensors")) - precond_path = tmp_path / "precond.safetensors" - save_file({MODULE: precond}, str(precond_path)) - return precond_path + preconditioner_path = tmp_path / "precond.safetensors" + save_file({MODULE: precond}, str(preconditioner_path)) + return preconditioner_path def _reference_diag_hessian(q_a, q_g, lam): @@ -98,16 +101,16 @@ def test_diagonal_preconditioner_matches_dense_reference(tmp_path, fn_kind): with diag(H) recovered exactly from the EKFAC factors.""" q_a, q_g, lam = _random_factors() precond = torch.rand(OUT_DIM, IN_DIM) + 0.5 - precond_path = _write_factor_shards(tmp_path, q_a, q_g, lam, precond) + preconditioner_path = _write_factor_shards(tmp_path, q_a, q_g, lam, precond) fn = {"f_backward": f_backward, "f_segment": f_segment}[fn_kind](LR_TIMES_STEPS) preconditioner = DiagonalFactoredPreconditioner.from_shards( tmp_path, - precond_path, + preconditioner_path, rank=0, device="cpu", apply_fn=fn, - multiply_by_precond=fn_kind == "f_segment", + multiply_by_preconditioner=fn_kind == "f_segment", ev_correction=True, ) @@ -131,15 +134,15 @@ def test_f_segment_zero_eigenvalue_limit(tmp_path): q_a, q_g, _ = _random_factors() lam = torch.zeros(OUT_DIM, IN_DIM) precond = torch.rand(OUT_DIM, IN_DIM) + 0.5 - precond_path = _write_factor_shards(tmp_path, q_a, q_g, lam, precond) + preconditioner_path = _write_factor_shards(tmp_path, q_a, q_g, lam, precond) preconditioner = DiagonalFactoredPreconditioner.from_shards( tmp_path, - precond_path, + preconditioner_path, rank=0, device="cpu", apply_fn=f_segment(LR_TIMES_STEPS), - multiply_by_precond=True, + multiply_by_preconditioner=True, ev_correction=True, ) grads = {MODULE: torch.ones(1, OUT_DIM * IN_DIM)} @@ -147,14 +150,14 @@ def test_f_segment_zero_eigenvalue_limit(tmp_path): torch.testing.assert_close(out, LR_TIMES_STEPS * precond) -def test_precond_shape_mismatch_raises(tmp_path): +def test_preconditioner_shape_mismatch_raises(tmp_path): q_a, q_g, lam = _random_factors() bad_precond = torch.rand(OUT_DIM + 1, IN_DIM) - precond_path = _write_factor_shards(tmp_path, q_a, q_g, lam, bad_precond) + preconditioner_path = _write_factor_shards(tmp_path, q_a, q_g, lam, bad_precond) with pytest.raises(ValueError, match="shape"): DiagonalFactoredPreconditioner.from_shards( tmp_path, - precond_path, + preconditioner_path, rank=0, device="cpu", apply_fn=f_backward(LR_TIMES_STEPS), @@ -162,7 +165,7 @@ def test_precond_shape_mismatch_raises(tmp_path): ) -def test_precond_missing_module_raises(tmp_path): +def test_preconditioner_missing_module_raises(tmp_path): q_a, q_g, lam = _random_factors() for sub, tensor in [ ("eigen_activation_sharded", q_a), @@ -171,12 +174,12 @@ def test_precond_missing_module_raises(tmp_path): ]: (tmp_path / sub).mkdir() save_file({MODULE: tensor}, str(tmp_path / sub / "shard_0.safetensors")) - precond_path = tmp_path / "precond.safetensors" - save_file({"other_module": torch.rand(OUT_DIM, IN_DIM)}, str(precond_path)) + preconditioner_path = tmp_path / "precond.safetensors" + save_file({"other_module": torch.rand(OUT_DIM, IN_DIM)}, str(preconditioner_path)) with pytest.raises(KeyError, match=MODULE): DiagonalFactoredPreconditioner.from_shards( tmp_path, - precond_path, + preconditioner_path, rank=0, device="cpu", apply_fn=f_backward(LR_TIMES_STEPS), @@ -184,6 +187,63 @@ def test_precond_missing_module_raises(tmp_path): ) +@pytest.mark.parametrize("fn_kind", ["f_backward", "f_segment"]) +def test_ekfac_applicator_uses_diagonal_path(tmp_path, fn_kind): + """``EkfacConfig.preconditioner_path`` routes ``apply_eigfn_to_query``'s + applicator through the diagonal preconditioner, and + ``preconditioner_post_multiply`` reaches ``multiply_by_preconditioner``.""" + q_a, q_g, lam = _random_factors() + precond = torch.rand(OUT_DIM, IN_DIM) + 0.5 + preconditioner_path = _write_factor_shards(tmp_path, q_a, q_g, lam, precond) + + query_path = tmp_path / "query" + index = create_index( + root=query_path, + num_grads=3, + grad_sizes={MODULE: OUT_DIM * IN_DIM}, + dtype=np.float32, + ) + index[:] = np.random.default_rng(0).standard_normal(index.shape).astype(np.float32) + index.flush() + + cfg = EkfacConfig( + hessian_method_path=str(tmp_path), + gradient_path=str(query_path), + run_path=str(tmp_path / "out"), + ev_correction=True, + preconditioner_path=str(preconditioner_path), + preconditioner_post_multiply=fn_kind == "f_segment", + ) + fn = {"f_backward": f_backward, "f_segment": f_segment}[fn_kind](LR_TIMES_STEPS) + EkfacApplicator(cfg, apply_fn=fn).compute_ivhp_sharded() + out = torch.from_numpy( + np.asarray(load_module_gradients(str(tmp_path / "out"))[MODULE][:]) + ) + + reference = DiagonalFactoredPreconditioner.from_shards( + tmp_path, + preconditioner_path, + rank=0, + device="cpu", + apply_fn=fn, + multiply_by_preconditioner=fn_kind == "f_segment", + ev_correction=True, + ).apply({MODULE: torch.from_numpy(np.asarray(index[:])).float()})[MODULE] + torch.testing.assert_close(out, reference.cpu()) + + +def test_ekfac_applicator_preconditioner_without_apply_fn_raises(tmp_path): + cfg = EkfacConfig( + hessian_method_path=str(tmp_path), + gradient_path=str(tmp_path), + run_path=str(tmp_path / "out"), + ev_correction=True, + preconditioner_path=str(tmp_path / "precond.safetensors"), + ) + with pytest.raises(ValueError, match="preconditioner_path requires apply_fn"): + EkfacApplicator(cfg).compute_ivhp_sharded() + + def test_build_segment_preconditioners(tmp_path): """Per-segment P built from the checkpoints' optimizer.pt files: bias-corrected via the stored step/betas, index-mapped to param names, From 0c81dc3910670ebfdc998081b731eb4701915dcb Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Sat, 25 Jul 2026 10:05:25 +0000 Subject: [PATCH 03/10] fix: don't scale gradients through to the caller in the diagonal apply DiagonalFactoredPreconditioner.apply reshaped the caller's tensor and scaled it in place. compute_ivhp_sharded passes gradients that alias a read-only mmap and relies on the preconditioner returning fresh tensors (as FactoredPreconditioner does, via its rotation matmuls). The aliasing only bites when the preceding `.to(device, float32)` is a no-op, i.e. on a CPU-only run: writing to the mapping segfaults. Copy before scaling. Caught by test_ekfac_applicator_uses_diagonal_path on CPU-only CI; also pin the no-write-through contract directly. --- bergson/hessians/preconditioner.py | 5 ++++- tests/test_source_optimizer_variants.py | 6 +++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/bergson/hessians/preconditioner.py b/bergson/hessians/preconditioner.py index ef81751e..5c0be5d4 100644 --- a/bergson/hessians/preconditioner.py +++ b/bergson/hessians/preconditioner.py @@ -450,7 +450,10 @@ def apply(self, grads: dict[str, Tensor]) -> dict[str, Tensor]: continue o = self.eigen_g[name].shape[1] i = self.eigen_a[name].shape[1] - g = flat.to(self.lambdas[name].device, torch.float32).view(-1, o, i) + # Copy rather than alias: the caller's gradients may be a read-only + # mmap, and the scale below writes through. `.to` is a no-op when + # they are already float32 on this device, as on a CPU-only run. + g = flat.to(self.lambdas[name].device, torch.float32).clone().view(-1, o, i) self.shard_computer.scale_rows_in_place(g, self._multiplier(name)) self.logger.debug("%s: scaled elementwise by f(p * diag(H))", name) out[name] = g.reshape(flat.shape[0], -1).to(flat.dtype) diff --git a/tests/test_source_optimizer_variants.py b/tests/test_source_optimizer_variants.py index 71866a23..5309e45f 100644 --- a/tests/test_source_optimizer_variants.py +++ b/tests/test_source_optimizer_variants.py @@ -115,7 +115,11 @@ def test_diagonal_preconditioner_matches_dense_reference(tmp_path, fn_kind): ) grads = {MODULE: torch.randn(3, OUT_DIM * IN_DIM)} - out = preconditioner.apply({k: v.clone() for k, v in grads.items()})[MODULE] + before = grads[MODULE].clone() + out = preconditioner.apply(grads)[MODULE] + # apply() must not write through to the caller's tensor: compute_ivhp_sharded + # hands it gradients aliasing a read-only mmap. + torch.testing.assert_close(grads[MODULE], before) sigma = precond * _reference_diag_hessian(q_a, q_g, lam) if fn_kind == "f_backward": From ec6abd1953f721df6af914e5684a1746a426770a Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Sat, 25 Jul 2026 10:24:30 +0000 Subject: [PATCH 04/10] perf: force the diagonal apply's copy with to(copy=True) The .to(...).clone() that fixed the read-only-mmap write copied twice on the GPU path: once host-to-device, once more on device. copy=True gets the same guarantee in a single copy. --- bergson/hessians/preconditioner.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/bergson/hessians/preconditioner.py b/bergson/hessians/preconditioner.py index 5c0be5d4..ccaffbd7 100644 --- a/bergson/hessians/preconditioner.py +++ b/bergson/hessians/preconditioner.py @@ -450,10 +450,12 @@ def apply(self, grads: dict[str, Tensor]) -> dict[str, Tensor]: continue o = self.eigen_g[name].shape[1] i = self.eigen_a[name].shape[1] - # Copy rather than alias: the caller's gradients may be a read-only - # mmap, and the scale below writes through. `.to` is a no-op when - # they are already float32 on this device, as on a CPU-only run. - g = flat.to(self.lambdas[name].device, torch.float32).clone().view(-1, o, i) + # copy=True because the scale below writes through: the caller's + # gradients may alias a read-only mmap, and a plain `.to` is a no-op + # when they are already float32 on this device, as on a CPU-only run. + g = flat.to(self.lambdas[name].device, torch.float32, copy=True).view( + -1, o, i + ) self.shard_computer.scale_rows_in_place(g, self._multiplier(name)) self.logger.debug("%s: scaled elementwise by f(p * diag(H))", name) out[name] = g.reshape(flat.shape[0], -1).to(flat.dtype) From 593039c5b44d6172124b513874b493815f0cc223 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 11:44:09 +0000 Subject: [PATCH 05/10] refactor: rename save_retrained_models to save_models The flag moved from ValidationConfig up to TrainingConfig, where "retrained" no longer describes what it does: on a plain training run nothing is re-trained, it just saves the model that run produced. Every config that carries it trains models, so name it for that. Metasmoothness gets the flag by inheritance rather than defining its own near-duplicate. Each run type still writes what it actually produces -- HF-format models under retrained/ for training and validation, theta_init.pt and theta_final.pt for metasmoothness -- and the docstring spells both out. The old key still parses via a deprecated save_retrained_models alias that warns and forwards, so existing YAML configs and model banks keep working; setting both to conflicting values is an error. The alias is marked for removal by 2026-10-29. --- bergson/cli/commands.py | 2 +- bergson/config/config.py | 36 ++++ bergson/magic/cli.py | 2 +- bergson/validate.py | 8 +- examples/bank_baselines/README.md | 41 +++++ examples/bank_baselines/__init__.py | 0 .../bank_baselines/activation_baseline.py | 153 ++++++++++++++++ examples/bank_baselines/bm25_baseline.py | 69 +++++++ examples/bank_baselines/common.py | 156 ++++++++++++++++ examples/bank_baselines/gradient_baseline.py | 92 ++++++++++ examples/bank_baselines/qwen3_baseline.py | 76 ++++++++ examples/bank_baselines/semantic_baseline.py | 102 +++++++++++ examples/magic/gpt2_wikitext.yaml | 2 +- examples/magic/gpt2_wikitext_bank.yaml | 43 +++++ examples/magic/gpt2_wikitext_dropout.yaml | 40 +++++ examples/magic/gpt2_wikitext_dropout_s15.yaml | 42 +++++ examples/magic/gpt2_wikitext_lotus.yaml | 44 +++++ .../magic/gpt2_wikitext_lotus_ekfac_q50.yaml | 79 ++++++++ .../gpt2_wikitext_lotus_final_q01_50.yaml | 31 ++++ examples/magic/gpt2_wikitext_lotus_mq.yaml | 70 ++++++++ .../gpt2_wikitext_train_epsroot0_bs64.yaml | 54 ++++++ .../gpt2_wikitext_train_stdadam_bs64.yaml | 54 ++++++ .../gpt2_wikitext_lotus_source_adam_q50.yaml | 91 ++++++++++ .../gpt2_wikitext_lotus_source_q50.yaml | 79 ++++++++ .../gpt2_wikitext_lotus_trackstar_q50.yaml | 170 ++++++++++++++++++ examples/pipelines/ekfac_qwen7b.yaml | 59 ++++++ examples/pipelines/recall_ekfac_olmo3.yaml | 63 +++++++ .../pipelines/recall_ekfac_olmo3_fp32.yaml | 62 +++++++ .../pipelines/recall_source_olmo3_fp32.yaml | 79 ++++++++ .../recall_trackstar_olmo3_fp32.yaml | 57 ++++++ .../recall_trackstar_olmo3_fp32_dim32.yaml | 57 ++++++ tests/test_banked_model_loading.py | 4 +- 32 files changed, 1908 insertions(+), 9 deletions(-) create mode 100644 examples/bank_baselines/README.md create mode 100644 examples/bank_baselines/__init__.py create mode 100644 examples/bank_baselines/activation_baseline.py create mode 100644 examples/bank_baselines/bm25_baseline.py create mode 100644 examples/bank_baselines/common.py create mode 100644 examples/bank_baselines/gradient_baseline.py create mode 100644 examples/bank_baselines/qwen3_baseline.py create mode 100644 examples/bank_baselines/semantic_baseline.py create mode 100644 examples/magic/gpt2_wikitext_bank.yaml create mode 100644 examples/magic/gpt2_wikitext_dropout.yaml create mode 100644 examples/magic/gpt2_wikitext_dropout_s15.yaml create mode 100644 examples/magic/gpt2_wikitext_lotus.yaml create mode 100644 examples/magic/gpt2_wikitext_lotus_ekfac_q50.yaml create mode 100644 examples/magic/gpt2_wikitext_lotus_final_q01_50.yaml create mode 100644 examples/magic/gpt2_wikitext_lotus_mq.yaml create mode 100644 examples/magic/gpt2_wikitext_train_epsroot0_bs64.yaml create mode 100644 examples/magic/gpt2_wikitext_train_stdadam_bs64.yaml create mode 100644 examples/method_comparison/gpt2_wikitext_lotus_source_adam_q50.yaml create mode 100644 examples/method_comparison/gpt2_wikitext_lotus_source_q50.yaml create mode 100644 examples/method_comparison/gpt2_wikitext_lotus_trackstar_q50.yaml create mode 100644 examples/pipelines/ekfac_qwen7b.yaml create mode 100644 examples/pipelines/recall_ekfac_olmo3.yaml create mode 100644 examples/pipelines/recall_ekfac_olmo3_fp32.yaml create mode 100644 examples/pipelines/recall_source_olmo3_fp32.yaml create mode 100644 examples/pipelines/recall_trackstar_olmo3_fp32.yaml create mode 100644 examples/pipelines/recall_trackstar_olmo3_fp32_dim32.yaml diff --git a/bergson/cli/commands.py b/bergson/cli/commands.py index a5f8bc89..d9e06765 100644 --- a/bergson/cli/commands.py +++ b/bergson/cli/commands.py @@ -276,7 +276,7 @@ class Validate(ValidationConfig): retrained_dir: str = "" """Optional: evaluate on an existing run directory of leave-k-out re-trained models written with - ``save_retrained_models=true``.""" + ``save_models=true``.""" def execute(self): """Run the validation.""" diff --git a/bergson/config/config.py b/bergson/config/config.py index 8d4cefa6..a1a1925e 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -1,5 +1,6 @@ import math import os +import warnings from abc import ABC from dataclasses import dataclass from pathlib import Path @@ -379,6 +380,26 @@ class TrainingConfig(AttributionConfig, Serializable): wandb_project: str = "" """Weights & Biases project name. If set, logs training loss to W&B.""" + save_models: bool = False + """When True, persist the models this run trains, in the form each run + produces them: + + - Training and validation runs save the fully-trained model (HF format, + weights + tokenizer) to ``/retrained/base/``, and validation + additionally saves each leave-k-out model to + ``/retrained/subset_/`` so it can be reused for later + attribution queries without retraining. ~0.5 GB per model for GPT-2. + - Metasmoothness runs save the unperturbed run's parameters to + ``/theta_final.pt`` and its initialization to ``theta_init.pt`` + (name -> fp32 CPU tensor, ~0.65 GB each for GPT-2). The two perturbed + runs are finite-difference probes and are not saved.""" + + # TODO(Lucia Quirke): remove by 2026-10-29. Backwards compatibility for the + # pre-rename key; the flag moved from ValidationConfig to TrainingConfig, + # where "retrained" no longer describes a plain training run. + save_retrained_models: bool | None = None + """Deprecated alias for ``save_models``.""" + def __post_init__(self): super().__post_init__() @@ -386,6 +407,21 @@ def __post_init__(self): if isinstance(self.save_optimizer_state, bool): self.save_optimizer_state = "last" if self.save_optimizer_state else "none" + # TODO(Lucia Quirke): remove by 2026-10-29, along with the field above. + if self.save_retrained_models is not None: + if self.save_models and self.save_models != self.save_retrained_models: + raise ValueError( + "save_retrained_models and save_models are set to conflicting " + "values; save_retrained_models is a deprecated alias, so set " + "only save_models." + ) + warnings.warn( + "save_retrained_models is deprecated; use save_models.", + FutureWarning, + stacklevel=2, + ) + self.save_models = self.save_retrained_models + @dataclass class MetasmoothnessConfig(TrainingConfig): diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index b3de45c3..c2041a75 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -284,7 +284,7 @@ def worker( os.path.join(run_cfg.run_path, "optimizer.pt"), ) - if run_cfg.save_retrained_models and global_rank == 0: + if run_cfg.save_models and global_rank == 0: # Persist the fully-trained (no leave-out) model alongside the per-subset # models, so a later evaluate_retrained run can measure the query # baseline from the bank itself instead of a separate base_model_dir. diff --git a/bergson/validate.py b/bergson/validate.py index 5621b266..447613d2 100644 --- a/bergson/validate.py +++ b/bergson/validate.py @@ -115,7 +115,7 @@ def bank_loss_cache_key( def _load_banked_model( run_cfg: ValidationConfig, out_dir: str, device: torch.device | str ) -> torch.nn.Module: - """Load a banked ``save_retrained_models`` checkpoint into a ready model.""" + """Load a banked ``save_models`` checkpoint into a ready model.""" load_kwargs = {"dtype": torch.float32, "attn_implementation": "eager"} load_kwargs.update(simple_parse_kwargs_string(run_cfg.model_kwargs)) if os.path.isfile(os.path.join(out_dir, "adapter_config.json")): @@ -330,7 +330,7 @@ def validate_scores( hf_set_verbosity_error() # Optionally persist each retrained model for later attribution queries. - save_models = getattr(run_cfg, "save_retrained_models", False) + save_models = getattr(run_cfg, "save_models", False) retrained_tokenizer = None if save_models and global_rank == 0: retrained_tokenizer = AutoTokenizer.from_pretrained( @@ -464,7 +464,7 @@ def evaluate_retrained( """Evaluate a bank of pre-saved leave-k-out models on a query, no retraining. Reads models written by an earlier run with - ``save_retrained_models=true`` and evaluates attribution scores. No training + ``save_models=true`` and evaluates attribution scores. No training happens so evaluation is cheap. """ assert score_path, "evaluate_retrained requires precomputed --scores" @@ -474,7 +474,7 @@ def evaluate_retrained( if not subsets_path.exists(): raise FileNotFoundError( f"{subsets_path} not found; retrained_dir must point at a run " - f"directory written with save_retrained_models=true" + f"directory written with save_models=true" ) run_path = Path(run_cfg.run_path) diff --git a/examples/bank_baselines/README.md b/examples/bank_baselines/README.md new file mode 100644 index 00000000..98765241 --- /dev/null +++ b/examples/bank_baselines/README.md @@ -0,0 +1,41 @@ +# Non-gradient attribution baselines + +Text similarity baselines evaluated by mean LDS over 50 test queries. + +- `bm25_baseline.py` — BM25 lexical overlap: pure surface-form term overlap, no model or embedding. +- `gradient_baseline.py` — gradient cosine similarity: cosine of the full per-example loss gradients on the bank's model (TracIn-style, unpreconditioned). +- `activation_baseline.py` — activation similarity: each doc is the mean-pooled input activation to every linear matrix of the model, L2-normalized per matrix and concatenated, then cosine similarity. +- `semantic_baseline.py` — semantic search with `jinaai/jina-embeddings-v3` (asymmetric `retrieval.query`/`retrieval.passage`). +- `qwen3_baseline.py` — semantic search with `Qwen/Qwen3-Embedding-8B`, a SOTA decoder embedder (`--model` to swap). + +Each produces a `[num_train_docs, num_queries]` score matrix. + +## Running + +Point a baseline at a re-train bank (written with `save_retrained_models=true`); it reads the model and dataset from the bank's `config.yaml`: + +```bash +python -m examples.bank_baselines.bm25_baseline --bank runs/retrain_bank_path +python -m examples.bank_baselines.gradient_baseline --bank runs/retrain_bank_path +python -m examples.bank_baselines.activation_baseline --bank runs/retrain_bank_path +python -m examples.bank_baselines.semantic_baseline --bank runs/retrain_bank_path +python -m examples.bank_baselines.qwen3_baseline --bank runs/retrain_bank_path +``` + +Omit `--bank` to build the default GPT-2/WikiText bank (`examples/magic/gpt2_wikitext_bank.yaml`) first. `--query_split` sets the query set (default `test[1:51]`), `--out` the output dir (default `runs/bank_baselines/`). + +## Results + +GPT-2 / WikiText bank (100 subsets, 1% leave-out, `eps_root 1e-8`): + +| Method | mean ρ | +| --- | --- | +| BM25 lexical overlap | 0.16 | +| Qwen3-Embedding-8B | 0.11 | +| activation similarity | 0.09 | +| Jina v3 semantic search | 0.06 | +| gradient cosine similarity | 0.05 | + +## Notes + +`jina-embeddings-v3`'s custom code predates transformers 5.x; `semantic_baseline.load_model` sets an `all_tied_weights_keys` default and resets the NaN LoRA `lora_dropout_mask` buffers to ones so it loads and runs. NVIDIA's NV-Embed-v2 is a comparable SOTA embedder but its custom code is incompatible with transformers 5.x, so `qwen3_baseline.py` uses Qwen3-Embedding instead. diff --git a/examples/bank_baselines/__init__.py b/examples/bank_baselines/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/bank_baselines/activation_baseline.py b/examples/bank_baselines/activation_baseline.py new file mode 100644 index 00000000..ac99fb0e --- /dev/null +++ b/examples/bank_baselines/activation_baseline.py @@ -0,0 +1,153 @@ +"""Baseline: activation similarity over the attribution target modules. + +Represents each document by the mean-pooled input activation to every target +linear module of the bank's fully-trained model (the same modules the gradient +methods attribute through: all Conv1D/Linear except ``lm_head``). Each +per-module pooled vector is L2-normalized, then concatenated, so cosine +similarity between two docs equals the average per-module cosine similarity. + +A training doc that is activation-similar to a query is predicted to be +influential (removing it should raise query loss), so the loss-diff-convention +score is ``-cosine`` -- larger => larger predicted loss reduction from keeping +the doc. + +Run with (builds the default bank if --bank is omitted): + python -m examples.bank_baselines.activation_baseline --bank runs/retrain_bank_path +""" + +import argparse +from pathlib import Path + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.pytorch_utils import Conv1D + +from . import common + + +def target_module_names(model: torch.nn.Module) -> list[str]: + """Linear/Conv1D modules the gradient methods attribute through (not lm_head).""" + names = [] + for name, mod in model.named_modules(): + if name.endswith("lm_head"): + continue + if isinstance(mod, (Conv1D, torch.nn.Linear)): + names.append(name) + return names + + +@torch.no_grad() +def embed( + texts: list[str], + model: torch.nn.Module, + tokenizer, + module_names: list[str], + device: str, + max_len: int, + batch_size: int, +) -> np.ndarray: + """Per-module-normalized, concatenated mean-pooled activation embeddings.""" + modules = dict(model.named_modules()) + captured: dict[str, torch.Tensor] = {} + handles = [] + for mn in module_names: + + def hook(_m, inp, _out, _mn=mn): + captured[_mn] = inp[0].detach() + + handles.append(modules[mn].register_forward_hook(hook)) + + out_rows = [] + try: + for start in range(0, len(texts), batch_size): + batch = texts[start : start + batch_size] + enc = tokenizer( + batch, + return_tensors="pt", + padding=True, + truncation=True, + max_length=max_len, + ).to(device) + mask = enc["attention_mask"].unsqueeze(-1).float() # [B, T, 1] + captured.clear() + model(**enc) + + per_module = [] + denom = mask.sum(dim=1).clamp_min(1.0) # [B, 1] + for mn in module_names: + act = captured[mn].float() # [B, T, d] + pooled = (act * mask).sum(dim=1) / denom # [B, d] + pooled = torch.nn.functional.normalize(pooled, dim=-1) + per_module.append(pooled) + emb = torch.cat(per_module, dim=-1) # [B, sum_d] + out_rows.append(emb.cpu().numpy()) + finally: + for h in handles: + h.remove() + return np.concatenate(out_rows, axis=0) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") + ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) + ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) + ap.add_argument("--device", default="cuda:0") + ap.add_argument("--max_len", type=int, default=1024) + ap.add_argument("--batch_size", type=int, default=32) + args = ap.parse_args() + + bank = common.ensure_bank(args.bank) + spec = common.read_bank_spec(bank) + out_dir = Path(args.out) + + tokenizer = AutoTokenizer.from_pretrained(spec.model) + tokenizer.pad_token = tokenizer.eos_token + model = AutoModelForCausalLM.from_pretrained( + spec.base_model, dtype=torch.float32, attn_implementation="eager" + ).to(args.device) + model.eval() + + module_names = target_module_names(model) + print(f"Embedding over {len(module_names)} modules, e.g. {module_names[:3]} ...") + + train_texts, query_texts = common.load_texts(spec, args.query_split) + print(f"Embedding {len(train_texts)} train docs ...") + train_emb = embed( + train_texts, + model, + tokenizer, + module_names, + args.device, + args.max_len, + args.batch_size, + ) + print(f"Embedding {len(query_texts)} query docs ...") + query_emb = embed( + query_texts, + model, + tokenizer, + module_names, + args.device, + args.max_len, + args.batch_size, + ) + + # Cosine similarity (rows already unit-norm per module; renormalize the + # concatenation so cosine == mean per-module cosine). + tn = train_emb / np.linalg.norm(train_emb, axis=1, keepdims=True) + qn = query_emb / np.linalg.norm(query_emb, axis=1, keepdims=True) + cosine = tn @ qn.T # [num_train, num_query] + + scores = -cosine # loss-diff convention (similar => influential) + score_path = common.save_scores(scores, out_dir, "activation_scores") + + rhos = common.evaluate_lds( + bank, score_path, out_dir / "activation_validate", spec, args.query_split + ) + common.report("activation similarity", rhos) + + +if __name__ == "__main__": + main() diff --git a/examples/bank_baselines/bm25_baseline.py b/examples/bank_baselines/bm25_baseline.py new file mode 100644 index 00000000..cf7f8928 --- /dev/null +++ b/examples/bank_baselines/bm25_baseline.py @@ -0,0 +1,69 @@ +"""Baseline: BM25 lexical overlap. + +Scores each training doc against each query by Okapi BM25 -- pure surface-form +term overlap, no model and no learned embedding. A lexical proxy for influence: +docs sharing query terms are predicted influential, so the loss-diff-convention +score is ``-bm25``. Often a stronger influence proxy than deep-semantic +embedders for small LMs, and essentially free. + +Run with (builds the default bank if --bank is omitted): + python -m examples.bank_baselines.bm25_baseline --bank runs/retrain_bank_path +""" + +import argparse +from pathlib import Path + +import numpy as np +from sklearn.feature_extraction.text import CountVectorizer + +from . import common + + +def bm25_scores( + train_texts: list[str], query_texts: list[str], k1: float = 1.5, b: float = 0.75 +) -> np.ndarray: + """Okapi BM25 score of every train doc against every query -> [docs, queries].""" + vectorizer = CountVectorizer() + tf = vectorizer.fit_transform(train_texts).astype(np.float64).tocsr() # [D, V] + n_docs, _ = tf.shape + + doc_len = np.asarray(tf.sum(axis=1)).ravel() + avgdl = doc_len.mean() + df = np.asarray((tf > 0).sum(axis=0)).ravel() + idf = np.log((n_docs - df + 0.5) / (df + 0.5) + 1.0) + + # Per-nonzero BM25 term weight: idf * tf*(k1+1) / (tf + k1*(1-b+b*dl/avgdl)). + rows, cols = tf.nonzero() + denom_bias = k1 * (1 - b + b * doc_len[rows] / avgdl) + weighted = tf.copy() + weighted.data = idf[cols] * tf.data * (k1 + 1.0) / (tf.data + denom_bias) + + # Query terms as a binary term set (Okapi BM25 ignores query term frequency). + q_bin = (vectorizer.transform(query_texts) > 0).astype(np.float64) # [Q, V] + return np.asarray((weighted @ q_bin.T).todense()) # [D, Q] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") + ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) + ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) + args = ap.parse_args() + + bank = common.ensure_bank(args.bank) + spec = common.read_bank_spec(bank) + out_dir = Path(args.out) + + train_texts, query_texts = common.load_texts(spec, args.query_split) + print(f"BM25 over {len(train_texts)} train docs, {len(query_texts)} queries ...") + scores = -bm25_scores(train_texts, query_texts) # loss-diff convention + score_path = common.save_scores(scores, out_dir, "bm25_scores") + + rhos = common.evaluate_lds( + bank, score_path, out_dir / "bm25_validate", spec, args.query_split + ) + common.report("BM25 lexical overlap", rhos) + + +if __name__ == "__main__": + main() diff --git a/examples/bank_baselines/common.py b/examples/bank_baselines/common.py new file mode 100644 index 00000000..2da7038b --- /dev/null +++ b/examples/bank_baselines/common.py @@ -0,0 +1,156 @@ +"""Shared helpers for non-gradient attribution baselines on a re-train bank. + +A "re-train bank" is any directory written with ``save_retrained_models=true`` +(e.g. by ``examples/magic/gpt2_wikitext_bank.yaml``): it holds ``subsets.json``, +``retrained/base`` and ``retrained/subset_*`` checkpoints, and the ``config.yaml`` +that built it. Each baseline reads the model + dataset straight from the bank, +produces a ``[num_train_docs, num_queries]`` score matrix, and evaluates it by +per-query LDS -- so ``--bank `` is all anyone needs to run one. + +The LDS itself reuses ``bergson``'s ``evaluate_retrained``: it computes each +subset's query-loss diff from the banked models once, caches it under the bank, +and every later baseline on the same bank/query reuses that cache. +""" + +import subprocess +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import yaml +from datasets import load_dataset + +from bergson.config.config import DataConfig +from bergson.magic.config import MagicConfig +from bergson.validate import evaluate_retrained + +REPO = Path(__file__).resolve().parents[2] + +# The default bank built on demand when ``--bank`` is omitted. +DEFAULT_BANK = REPO / "runs" / "gpt2_wikitext_bank" +DEFAULT_BANK_CONFIG = REPO / "examples" / "magic" / "gpt2_wikitext_bank.yaml" +# LDS query set (50 held-out docs); overridable per run. +DEFAULT_QUERY_SPLIT = "test[1:51]" + + +@dataclass +class BankSpec: + """The model + data a bank was built on, read from its ``config.yaml``.""" + + model: str + base_model: str # retrained/base if present, else the base model id + dataset: str + train_split: str + prompt_column: str + chunk_length: int + batch_size: int + + +def ensure_bank(bank: str | None) -> Path: + """Return a ready re-train bank, building the default one if none is given. + + A bank is ready when it has ``subsets.json`` and ``retrained/base``. When + ``bank`` is ``None`` and the default bank is absent, it is built by running + its config through ``bergson`` (a leave-k-out re-training job). + """ + if bank is not None: + path = Path(bank) + if not (path / "subsets.json").exists(): + raise FileNotFoundError( + f"{path} is not a re-train bank (no subsets.json); pass a dir " + "written with save_retrained_models=true" + ) + return path + + if not (DEFAULT_BANK / "subsets.json").exists(): + cmd = ["bergson", str(DEFAULT_BANK_CONFIG)] + print( + f"No bank passed and {DEFAULT_BANK} missing; building it:\n" + f" {' '.join(cmd)}" + ) + subprocess.run(cmd, cwd=REPO, check=True) + return DEFAULT_BANK + + +def read_bank_spec(bank: Path) -> BankSpec: + """Read the model + dataset a bank was built on from its ``config.yaml``.""" + with open(bank / "config.yaml") as f: + config = yaml.safe_load(f) + # The bank-building step (magic/train) is the one carrying model + data. + step = next( + v + for s in config["steps"] + for v in s.values() + if isinstance(v, dict) and "model" in v and "data" in v + ) + data = step["data"] + base = bank / "retrained" / "base" + return BankSpec( + model=step["model"], + base_model=str(base) if base.exists() else step["model"], + dataset=data["dataset"], + train_split=data.get("split", "train"), + prompt_column=data.get("prompt_column", "text"), + chunk_length=data.get("chunk_length", 0), + batch_size=step.get("batch_size", 64), + ) + + +def load_texts(spec: BankSpec, query_split: str) -> tuple[list[str], list[str]]: + """Return (train_texts, query_texts) in original doc order. + + Row i of ``train_texts`` is training doc id i, matching the ids in the + bank's ``subsets.json`` and the rows of every score matrix. + """ + train = load_dataset(spec.dataset, split=spec.train_split) + query = load_dataset(spec.dataset, split=query_split) + return list(train[spec.prompt_column]), list(query[spec.prompt_column]) + + +def evaluate_lds( + bank: Path, + score_path: Path, + out_dir: Path, + spec: BankSpec, + query_split: str, +) -> np.ndarray: + """Per-query LDS Spearman of a score matrix against the bank. + + Runs ``evaluate_retrained`` (no re-training; reuses the bank's cached + per-subset query losses) and returns the per-query Spearman correlations it + writes to ``summary.csv``. + """ + run_cfg = MagicConfig( + run_path=str(out_dir), + model=spec.model, + precision="fp32", + batch_size=spec.batch_size, + query=DataConfig( + dataset=spec.dataset, + split=query_split, + prompt_column=spec.prompt_column, + chunk_length=0, + ), + ) + evaluate_retrained(run_cfg, str(bank), score_path=str(score_path)) + + summary = np.genfromtxt(out_dir / "summary.csv", delimiter=",", names=True) + return np.atleast_1d(summary["spearman_corr"]).astype(float) + + +def save_scores(scores: np.ndarray, out_dir: Path, name: str) -> Path: + """Save a ``[num_train_docs, num_queries]`` score matrix as ``.npy``.""" + out_dir.mkdir(parents=True, exist_ok=True) + path = out_dir / f"{name}.npy" + np.save(path, scores.astype(np.float32)) + print(f"Saved scores {scores.shape} -> {path}") + return path + + +def report(name: str, rhos: np.ndarray) -> None: + """Print a per-query LDS summary consistent with the repo's convention.""" + print(f"\n=== {name}: per-query LDS Spearman (n={len(rhos)} queries) ===") + print(f" mean {np.mean(rhos):+.4f}") + print(f" median {np.median(rhos):+.4f}") + print(f" min {np.min(rhos):+.4f} max {np.max(rhos):+.4f}") + print(f" frac>0 {np.mean(rhos > 0):.2f}") diff --git a/examples/bank_baselines/gradient_baseline.py b/examples/bank_baselines/gradient_baseline.py new file mode 100644 index 00000000..23416113 --- /dev/null +++ b/examples/bank_baselines/gradient_baseline.py @@ -0,0 +1,92 @@ +"""Baseline: gradient cosine similarity (TracIn-style, full parameters). + +Scores a train doc against a query by the cosine similarity of their full +per-example loss gradients on the bank's model -- the raw, unpreconditioned +influence signal that TrackStar/EK-FAC refine. A train doc whose gradient +aligns with the query's is predicted influential, so the loss-diff-convention +score is ``-cosine``. + +Full gradients are never stored: the 50 query gradients stay resident on the +GPU (GPT-2 is 124M params, ~25 GB in fp32) and each train-doc gradient is +computed, dotted against all queries, and discarded. + +Run with (builds the default bank if --bank is omitted): + python -m examples.bank_baselines.gradient_baseline --bank runs/retrain_bank_path +""" + +import argparse +from pathlib import Path + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from . import common + + +def per_example_grad( + text: str, model, tokenizer, device: str, max_len: int +) -> torch.Tensor: + """Flattened full-parameter gradient of the mean-CE loss on one document.""" + enc = tokenizer(text, return_tensors="pt", truncation=True, max_length=max_len).to( + device + ) + model.zero_grad(set_to_none=True) + model(input_ids=enc["input_ids"], labels=enc["input_ids"]).loss.backward() + return torch.cat( + [ + (p.grad if p.grad is not None else torch.zeros_like(p)).flatten() + for p in model.parameters() + ] + ) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") + ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) + ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) + ap.add_argument("--device", default="cuda:0") + ap.add_argument("--max_len", type=int, default=1024) + args = ap.parse_args() + + bank = common.ensure_bank(args.bank) + spec = common.read_bank_spec(bank) + out_dir = Path(args.out) + + tokenizer = AutoTokenizer.from_pretrained(spec.model) + model = AutoModelForCausalLM.from_pretrained( + spec.base_model, dtype=torch.float32, attn_implementation="eager" + ).to(args.device) + model.eval() + + train_texts, query_texts = common.load_texts(spec, args.query_split) + + # Keep the query gradients resident; stream train gradients against them. + # Pre-allocate and fill row by row so we never hold a second copy (each full + # gradient is ~0.5 GB, the [Q, P] block ~25 GB in fp32). + print(f"Computing {len(query_texts)} query gradients ...") + n_params = sum(p.numel() for p in model.parameters()) + query_grads = torch.empty(len(query_texts), n_params, device=args.device) + for i, t in enumerate(query_texts): + query_grads[i] = per_example_grad(t, model, tokenizer, args.device, args.max_len) + query_norms = query_grads.norm(dim=1).clamp_min(1e-12) + + print(f"Scoring {len(train_texts)} train gradients against queries ...") + scores = np.empty((len(train_texts), len(query_texts)), dtype=np.float32) + for i, text in enumerate(train_texts): + g = per_example_grad(text, model, tokenizer, args.device, args.max_len) + cos = (query_grads @ g) / (query_norms * g.norm().clamp_min(1e-12)) + scores[i] = cos.detach().cpu().numpy() + + scores = -scores # loss-diff convention (aligned gradient => influential) + score_path = common.save_scores(scores, out_dir, "gradient_scores") + + rhos = common.evaluate_lds( + bank, score_path, out_dir / "gradient_validate", spec, args.query_split + ) + common.report("gradient cosine similarity", rhos) + + +if __name__ == "__main__": + main() diff --git a/examples/bank_baselines/qwen3_baseline.py b/examples/bank_baselines/qwen3_baseline.py new file mode 100644 index 00000000..42576ace --- /dev/null +++ b/examples/bank_baselines/qwen3_baseline.py @@ -0,0 +1,76 @@ +"""Baseline: semantic search with Qwen3-Embedding-8B. + +A SOTA decoder-based embedding model (causal-attention Qwen3 backbone, +last-token pooling), top of the open MTEB leaderboard -- a much stronger, +attention-based embedder than the jina-v3 encoder. Training docs are embedded +as passages and queries with the retrieval prompt; a train doc scores against a +query by cosine similarity, so the loss-diff-convention score is ``-cosine``. + +(NVIDIA's NV-Embed-v2 is a comparable model but its custom modeling code is +incompatible with transformers 5.x; Qwen3-Embedding uses the native Qwen3 +architecture and loads cleanly. Pass --model Qwen/Qwen3-Embedding-4B for the +smaller, faster variant.) + +Run with (builds the default bank if --bank is omitted): + python -m examples.bank_baselines.qwen3_baseline --bank runs/retrain_bank_path +""" + +import argparse +from pathlib import Path + +from sentence_transformers import SentenceTransformer + +from . import common + +MODEL = "Qwen/Qwen3-Embedding-8B" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") + ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) + ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) + ap.add_argument("--model", default=MODEL) + ap.add_argument("--device", default="cuda:0") + ap.add_argument("--max_length", type=int, default=512) + ap.add_argument("--batch_size", type=int, default=8) + args = ap.parse_args() + + bank = common.ensure_bank(args.bank) + spec = common.read_bank_spec(bank) + out_dir = Path(args.out) + + model = SentenceTransformer( + args.model, model_kwargs={"torch_dtype": "bfloat16"}, device=args.device + ) + model.max_seq_length = args.max_length + + train_texts, query_texts = common.load_texts(spec, args.query_split) + print(f"Embedding {len(train_texts)} train docs (passages) ...") + train_emb = model.encode( + train_texts, + batch_size=args.batch_size, + normalize_embeddings=True, + show_progress_bar=True, + ) + print(f"Embedding {len(query_texts)} query docs (query prompt) ...") + query_emb = model.encode( + query_texts, + prompt_name="query", + batch_size=args.batch_size, + normalize_embeddings=True, + show_progress_bar=True, + ) + + cosine = train_emb @ query_emb.T # rows unit-norm => dot == cosine + scores = -cosine # loss-diff convention (similar => influential) + score_path = common.save_scores(scores, out_dir, "qwen3_scores") + + rhos = common.evaluate_lds( + bank, score_path, out_dir / "qwen3_validate", spec, args.query_split + ) + common.report(f"{args.model} semantic similarity", rhos) + + +if __name__ == "__main__": + main() diff --git a/examples/bank_baselines/semantic_baseline.py b/examples/bank_baselines/semantic_baseline.py new file mode 100644 index 00000000..2103251d --- /dev/null +++ b/examples/bank_baselines/semantic_baseline.py @@ -0,0 +1,102 @@ +"""Baseline: semantic-search similarity with a Jina AI embedding model. + +Embeds each training document as a retrieval passage and each query document +as a retrieval query with ``jinaai/jina-embeddings-v3`` (a strong 570M-param, +1024-dim, 8192-context text embedder), then scores a train doc against a query +by their cosine similarity -- ordinary dense semantic search. This ignores the +attributed model entirely; it is a pure content-similarity baseline. + +A training doc semantically similar to a query is predicted to be influential +(removing it should raise query loss), so the loss-diff-convention score is +``-cosine``. + +jina-embeddings-v3 ships custom modeling code that predates transformers 5.x, +so ``load_model`` patches two load-time incompatibilities (see there) to run it +for inference. + +Run with (builds the default bank if --bank is omitted): + python -m examples.bank_baselines.semantic_baseline --bank runs/retrain_bank_path +""" + +import argparse +from pathlib import Path + +import numpy as np +import torch +from transformers.modeling_utils import PreTrainedModel + +from . import common + +MODEL = "jinaai/jina-embeddings-v3" + + +def load_model(device: str): + # jina-embeddings-v3's custom code predates transformers 5.x; two fixes: + # (1) from_pretrained reads all_tied_weights_keys, which the custom class + # never defines -- give a benign empty default so load doesn't crash. + if not isinstance( + getattr(PreTrainedModel, "all_tied_weights_keys", None), property + ): + PreTrainedModel.all_tied_weights_keys = {} + from transformers import AutoModel + + model = AutoModel.from_pretrained( + MODEL, trust_remote_code=True, dtype=torch.float32 + ) + + # (2) its LoRA task adapters leave the per-forward lora_dropout_mask buffers + # uninitialized (NaN), so every embedding comes out NaN. In eval there + # is no dropout, so reset them to ones. + for name, buf in model.named_buffers(): + if "lora_dropout_mask" in name: + buf.data = torch.ones_like(buf) + return model.to(device).eval() + + +@torch.no_grad() +def encode(model, texts: list[str], task: str, batch_size: int) -> np.ndarray: + """L2-normalized embeddings via jina's task-specific encode API.""" + out = [] + for start in range(0, len(texts), batch_size): + emb = model.encode( + texts[start : start + batch_size], + task=task, + convert_to_numpy=True, + normalize_embeddings=True, + ) + out.append(np.asarray(emb, dtype=np.float32)) + return np.concatenate(out, axis=0) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") + ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) + ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) + ap.add_argument("--device", default="cuda:0") + ap.add_argument("--batch_size", type=int, default=16) + args = ap.parse_args() + + bank = common.ensure_bank(args.bank) + spec = common.read_bank_spec(bank) + out_dir = Path(args.out) + model = load_model(args.device) + + train_texts, query_texts = common.load_texts(spec, args.query_split) + print(f"Embedding {len(train_texts)} train docs (retrieval.passage) ...") + train_emb = encode(model, train_texts, "retrieval.passage", args.batch_size) + print(f"Embedding {len(query_texts)} query docs (retrieval.query) ...") + query_emb = encode(model, query_texts, "retrieval.query", args.batch_size) + + cosine = train_emb @ query_emb.T # rows unit-norm => dot == cosine + scores = -cosine # loss-diff convention (similar => influential) + score_path = common.save_scores(scores, out_dir, "semantic_scores") + + rhos = common.evaluate_lds( + bank, score_path, out_dir / "semantic_validate", spec, args.query_split + ) + common.report("Jina v3 semantic similarity", rhos) + + +if __name__ == "__main__": + main() diff --git a/examples/magic/gpt2_wikitext.yaml b/examples/magic/gpt2_wikitext.yaml index 29ac6a30..7cfc4649 100644 --- a/examples/magic/gpt2_wikitext.yaml +++ b/examples/magic/gpt2_wikitext.yaml @@ -49,4 +49,4 @@ steps: wandb_project: magic save_optimizer_state: True - save_retrained_models: True + save_models: True diff --git a/examples/magic/gpt2_wikitext_bank.yaml b/examples/magic/gpt2_wikitext_bank.yaml new file mode 100644 index 00000000..f47758b4 --- /dev/null +++ b/examples/magic/gpt2_wikitext_bank.yaml @@ -0,0 +1,43 @@ +# Build a GPT-2 / WikiText leave-k-out re-train bank (single node): MAGIC +# attribution plus 400 retrained leave-out models saved for later evaluation. + +# Run with `bergson examples/magic/gpt2_wikitext_bank.yaml` + +steps: + - magic: + run_path: runs/gpt2_wikitext_bank + model: gpt2 + overwrite: true + cleanup_ckpts: false + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:1]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + eps_root: 1e-8 + + batch_size: 64 + num_epochs: 4 + lr_schedule: + lr_scheduler_type: polynomial + lr: 0.0008 + lr_start: 1e-6 + lr_end: 0.00008 + warmup_steps: 0.25 + + subset_strategy: random + wandb_project: magic + num_subsets: 100 + subset_fraction: 0.01 + save_optimizer_state: True + save_retrained_models: True diff --git a/examples/magic/gpt2_wikitext_dropout.yaml b/examples/magic/gpt2_wikitext_dropout.yaml new file mode 100644 index 00000000..6c3c04a0 --- /dev/null +++ b/examples/magic/gpt2_wikitext_dropout.yaml @@ -0,0 +1,40 @@ +# LDS: 0.1862 + +steps: + - magic: + run_path: runs/gpt2_wikitext_dropout + model: gpt2 + overwrite: true + cleanup_ckpts: false + train_mode: true + weight_decay: 0.01 + adam_eps: 1.0e-8 + eps_root: 1.0e-6 + loss_reduction: mean + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:1]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + batch_size: 64 + num_epochs: 4 + lr_schedule: + lr_scheduler_type: polynomial + lr: 0.0008 + lr_start: 1e-6 + lr_end: 0.00008 + warmup_steps: 0.25 + + wandb_project: magic + save_optimizer_state: True + save_retrained_models: True diff --git a/examples/magic/gpt2_wikitext_dropout_s15.yaml b/examples/magic/gpt2_wikitext_dropout_s15.yaml new file mode 100644 index 00000000..a9e7cd8c --- /dev/null +++ b/examples/magic/gpt2_wikitext_dropout_s15.yaml @@ -0,0 +1,42 @@ +# LDS: -0.2286 + +steps: + - magic: + run_path: runs/gpt2_wikitext_dropout_s15 + model: gpt2 + overwrite: true + cleanup_ckpts: false + train_mode: true + num_subsets: 15 + subset_fraction: 0.01 + weight_decay: 0.01 + adam_eps: 1.0e-8 + eps_root: 1.0e-6 + loss_reduction: mean + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:1]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + batch_size: 64 + num_epochs: 4 + lr_schedule: + lr_scheduler_type: polynomial + lr: 0.0008 + lr_start: 1e-6 + lr_end: 0.00008 + warmup_steps: 0.25 + + wandb_project: magic + save_optimizer_state: True + save_retrained_models: True diff --git a/examples/magic/gpt2_wikitext_lotus.yaml b/examples/magic/gpt2_wikitext_lotus.yaml new file mode 100644 index 00000000..e1df5f3b --- /dev/null +++ b/examples/magic/gpt2_wikitext_lotus.yaml @@ -0,0 +1,44 @@ +# Single-step MAGIC attribution example (GPT-2 / WikiText, single node). +# Isolated run_path ("lotus") so this doesn't collide with any other +# job sharing examples/magic/gpt2_wikitext.yaml's run_path. + +# Run with `bergson examples/magic/gpt2_wikitext_lotus.yaml` + +steps: + - magic: + run_path: runs/lotus + model: gpt2 + overwrite: true + cleanup_ckpts: false + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:1]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + eps_root: 1e-6 + + batch_size: 64 + num_epochs: 4 + lr_schedule: + lr_scheduler_type: polynomial + lr: 0.0008 + lr_start: 1e-6 + lr_end: 0.00008 + warmup_steps: 0.25 + + subset_strategy: random + wandb_project: magic + num_subsets: 400 + subset_fraction: 0.01 + save_optimizer_state: True + save_retrained_models: True diff --git a/examples/magic/gpt2_wikitext_lotus_ekfac_q50.yaml b/examples/magic/gpt2_wikitext_lotus_ekfac_q50.yaml new file mode 100644 index 00000000..e84af898 --- /dev/null +++ b/examples/magic/gpt2_wikitext_lotus_ekfac_q50.yaml @@ -0,0 +1,79 @@ +# EK-FAC counterpart to gpt2_wikitext_lotus_mq.yaml: computes EK-FAC +# attribution scores for the SAME base model trained by +# examples/magic/gpt2_wikitext_lotus.yaml (runs/lotus/retrained/base, +# bs=64/eps_root=1e-6), against the same 50 test queries (test[1:51]), +# then validates those scores against the existing 400-model leave-k-out +# bank in runs/lotus/retrained -- no retraining needed, that bank is +# already complete. +# +# IMPORTANT: run_path values below are meant to be unique on this shared +# filesystem. If another run is already using them, rename before running +# (e.g. add your hostname) -- do NOT run this against a run_path someone +# else might be using. runs/lotus itself is read-only here and already +# finished, so it's safe to read from concurrently. + +# Run with `bergson examples/magic/gpt2_wikitext_lotus_ekfac_q50.yaml` +# from /mnt/ssd-1/lucia/bergson-damping (paths below are relative to that). + +steps: + - ekfac: + index_cfg: + run_path: runs/lotus_scores_ekfac50q + model: runs/lotus/retrained/base + overwrite: true + # GPT-2's max context length is 1024. + token_batch_size: 1024 + # Cap batch size in docs so the bin-packer doesn't produce + # large batches of very few tokens (gradients aren't compressed + # for EK-FAC). + max_batch_size: 64 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 # <-- adjust to this node's GPU count + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + hessian_cfg: + method: kfac + ev_correction: true + + score_cfg: + batch_size: 1024 + + preprocess_cfg: + unit_normalize: false + + hessian_pipeline_cfg: + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + # One score column per query (50 total), not aggregated -- + # required for the multi-query LDS validation below. + query_aggregation: none + inversion_cfg: + damping_factor: 0.1 + + - validate: + run_path: runs/lotus_scores_ekfac50q_validate + model: gpt2 + overwrite: true + + scores: runs/lotus_scores_ekfac50q/scores + retrained_dir: runs/lotus + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + + batch_size: 64 diff --git a/examples/magic/gpt2_wikitext_lotus_final_q01_50.yaml b/examples/magic/gpt2_wikitext_lotus_final_q01_50.yaml new file mode 100644 index 00000000..26767b82 --- /dev/null +++ b/examples/magic/gpt2_wikitext_lotus_final_q01_50.yaml @@ -0,0 +1,31 @@ +# Final true multi-query LDS check on all 50 per-query MAGIC backward +# passes, recomputed post-fix by scratchpad/run_50_bwd_fixed.sh. No +# retraining, no backward pass -- reuses the existing 400-model leave-k-out +# bank at runs/lotus/retrained, scored against a stacked [4608, 50] +# multi-query score array (runs/lotus_bwd/final_q01_50_scores.npy). +# +# Run with `bergson examples/magic/gpt2_wikitext_lotus_final_q01_50.yaml` + +steps: + - validate: + run_path: runs/lotus_final_q01_50 + model: gpt2 + overwrite: true + + scores: runs/lotus_bwd/final_q01_50_scores.npy + retrained_dir: runs/lotus + + distributed: + nproc_per_node: 1 + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + + batch_size: 8 diff --git a/examples/magic/gpt2_wikitext_lotus_mq.yaml b/examples/magic/gpt2_wikitext_lotus_mq.yaml new file mode 100644 index 00000000..fe9f2bcb --- /dev/null +++ b/examples/magic/gpt2_wikitext_lotus_mq.yaml @@ -0,0 +1,70 @@ +# Multi-query LDS follow-up to the "lotus" run (examples/magic/gpt2_wikitext_lotus.yaml). +# +# Step 1 retrains the same bs=64/eps_root=1e-6 base model (deterministic, +# same seed) purely to get a fresh MAGIC backward pass against 50 NEW test +# queries (test[1:51], i.e. the 50 WikiText test docs after the one used by +# the original lotus run) -- skip_validation stops it before any retraining. +# +# Step 2 reuses the 400 leave-k-out checkpoints already saved under +# runs/lotus/retrained (no retraining) and correlates them against the new +# per-query scores from step 1, reporting a mean Spearman LDS across the +# 50 queries. + +# Run with `bergson examples/magic/gpt2_wikitext_lotus_mq.yaml` + +steps: + - magic: + run_path: runs/lotus_mq_scores + model: gpt2 + overwrite: true + cleanup_ckpts: false + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + eps_root: 1e-6 + + batch_size: 64 + num_epochs: 4 + lr_schedule: + lr_scheduler_type: polynomial + lr: 0.0008 + lr_start: 1e-6 + lr_end: 0.00008 + warmup_steps: 0.25 + + skip_validation: true + save_retrained_models: false + save_optimizer_state: false + wandb_project: magic + + - validate: + run_path: runs/lotus_mq_eval + model: gpt2 + overwrite: true + + scores: runs/lotus_mq_scores/scores.pt + retrained_dir: runs/lotus + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + + batch_size: 64 diff --git a/examples/magic/gpt2_wikitext_train_epsroot0_bs64.yaml b/examples/magic/gpt2_wikitext_train_epsroot0_bs64.yaml new file mode 100644 index 00000000..9aa8e6c0 --- /dev/null +++ b/examples/magic/gpt2_wikitext_train_epsroot0_bs64.yaml @@ -0,0 +1,54 @@ +steps: +- train: + run_path: runs/gpt2_wikitext_epsroot0_bs64 + overwrite: true + model: gpt2 + precision: fp32 + revision: null + distributed: + nnode: 1 + nproc_per_node: 8 + node_rank: null + fsdp: false + peft_init_kwargs: '' + model_kwargs: '' + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: train + subset: null + prompt_column: text + completion_column: '' + conversation_column: '' + reward_column: '' + skip_nan_rewards: false + truncation: false + format_template: '' + data_kwargs: '' + chunk_length: 0 + tokenizer: '' + drop_columns: true + max_tokens: null + use_tf32_matmuls: false + debug: false + lr_schedule: + lr: 0.0008 + lr_scheduler_type: polynomial + lr_start: 1.0e-06 + lr_end: 8.0e-05 + warmup_steps: 0.25 + num_cycles: 0.5 + power: 1.0 + batch_size: 64 + num_epochs: 4 + seed: 42 + adam_beta1: 0.95 + adam_beta2: 0.975 + eps_root: 0.0 + optimizer: adamw + save_optimizer_state: true + save_retrained_models: true + weight_decay: 0.01 + max_grad_norm: null + grad_checkpointing: false + resume: false + wandb_project: magic diff --git a/examples/magic/gpt2_wikitext_train_stdadam_bs64.yaml b/examples/magic/gpt2_wikitext_train_stdadam_bs64.yaml new file mode 100644 index 00000000..45abebb9 --- /dev/null +++ b/examples/magic/gpt2_wikitext_train_stdadam_bs64.yaml @@ -0,0 +1,54 @@ +steps: +- train: + run_path: runs/gpt2_wikitext_stdadam_bs64 + overwrite: true + model: gpt2 + precision: fp32 + revision: null + distributed: + nnode: 1 + nproc_per_node: 8 + node_rank: null + fsdp: false + peft_init_kwargs: '' + model_kwargs: '' + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: train + subset: null + prompt_column: text + completion_column: '' + conversation_column: '' + reward_column: '' + skip_nan_rewards: false + truncation: false + format_template: '' + data_kwargs: '' + chunk_length: 0 + tokenizer: '' + drop_columns: true + max_tokens: null + use_tf32_matmuls: false + debug: false + lr_schedule: + lr: 0.0008 + lr_scheduler_type: polynomial + lr_start: 1.0e-06 + lr_end: 8.0e-05 + warmup_steps: 0.25 + num_cycles: 0.5 + power: 1.0 + batch_size: 64 + num_epochs: 4 + seed: 42 + adam_beta1: 0.9 + adam_beta2: 0.999 + eps_root: 0.0 + optimizer: adamw + save_optimizer_state: true + save_retrained_models: true + weight_decay: 0.01 + max_grad_norm: null + grad_checkpointing: false + resume: false + wandb_project: magic diff --git a/examples/method_comparison/gpt2_wikitext_lotus_source_adam_q50.yaml b/examples/method_comparison/gpt2_wikitext_lotus_source_adam_q50.yaml new file mode 100644 index 00000000..467c03a1 --- /dev/null +++ b/examples/method_comparison/gpt2_wikitext_lotus_source_adam_q50.yaml @@ -0,0 +1,91 @@ +# Adam/AdamW SOURCE variant (Bae et al., 2024, Appendix C) for the lotus +# GPT-2 base model against the 50 test queries, validated against the +# complete 400-model lotus bank. The unrolling eigenfunctions are evaluated +# on a diagonal approximation of P^1/2 H P^1/2, where P is the per-segment +# diagonal Adam preconditioner built from each checkpoint dir's +# optimizer.pt (written per snapshot during training by +# TrainingConfig.save_optimizer_state and copied in at checkpoint export). +# +# Reuses the lotus_source_q50 run's query gradients and segment KFAC/lambda +# artifacts via symlinks; only the eigenfunction application (steps 6-7) and +# scoring (step 8) recompute. Same 6-checkpoint/3-segment recipe and lr_list. +# +# Validated LDS at N=400: mean Spearman rho=0.207 across the 50 queries -- +# weaker than the SGD variant's 0.390 on the same bank, suggesting the +# paper's diagonal-Hessian approximation costs more here than the +# optimizer-correct preconditioning gains. + +# Run with `bergson examples/method_comparison/gpt2_wikitext_lotus_source_adam_q50.yaml` + +steps: + - approxunrolling: + index_cfg: + run_path: runs/lotus_source_adam_q50 + # Tokenizer + final-model reference; per-checkpoint passes load the + # model dirs below. + model: runs/lotus/retrained/base + # In the approxunrolling pipeline `overwrite` doubles as resume: + # steps whose output dirs exist are skipped after a crash. + overwrite: true + precision: fp32 + # GPT-2's max context length is 1024. + token_batch_size: 1024 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 # <-- adjust to this node's GPU count + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + # Doc-level scoring (one doc per row, no packing): packed 512-token + # chunks contaminate per-doc gradients across doc boundaries and + # were found to suppress LDS (see gpt2_lotus_trackstar50q_docs). + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + approx_unrolling_cfg: + checkpoints: + - runs/lotus_source_q50/models/step_51 + - runs/lotus_source_q50/models/step_102 + - runs/lotus_source_q50/models/step_136 + - runs/lotus_source_q50/models/step_187 + - runs/lotus_source_q50/models/step_238 + - runs/lotus_source_q50/models/step_288 + segments: 3 + lr_list: [4.857423e-04, 5.616667e-04, 2.416667e-04] + step_size_list: [97, 96, 96] + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + # One score column per query (50 total), like the EK-FAC and + # TrackStar runs, for multi-query LDS validation. + query_aggregation: none + # Adam variant: build per-segment diagonal preconditioners from the + # checkpoints' optimizer.pt files. + use_adam_preconditioner: true + + - validate: + run_path: runs/lotus_source_adam_q50_validate + model: gpt2 + overwrite: true + + scores: runs/lotus_source_adam_q50/scores.npy + retrained_dir: runs/lotus + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 512 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + + batch_size: 64 diff --git a/examples/method_comparison/gpt2_wikitext_lotus_source_q50.yaml b/examples/method_comparison/gpt2_wikitext_lotus_source_q50.yaml new file mode 100644 index 00000000..c7cafaba --- /dev/null +++ b/examples/method_comparison/gpt2_wikitext_lotus_source_q50.yaml @@ -0,0 +1,79 @@ +# SOURCE (approximate unrolling, Bae et al. 2024) attribution scores for the +# lotus GPT-2 base model against the 50 test queries, validated against the +# complete 400-model lotus bank -- the setting where MAGIC scores rho=0.92. +# Plain SGD variant (no optimizer preconditioning); see +# gpt2_wikitext_lotus_source_adam_q50.yaml for the Adam/AdamW variant. +# +# Validated LDS at N=400: mean Spearman rho=0.390 across the 50 queries. + +# Run with `bergson examples/method_comparison/gpt2_wikitext_lotus_source_q50.yaml` + +steps: + - approxunrolling: + index_cfg: + run_path: runs/lotus_source_q50_damp0 + # Tokenizer + final-model reference; per-checkpoint passes load the + # model dirs below. + model: runs/lotus/retrained/base + # In the approxunrolling pipeline `overwrite` doubles as resume: + # steps whose output dirs exist are skipped after a crash. + overwrite: true + precision: fp32 + # GPT-2's max context length is 1024. + token_batch_size: 1024 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 # <-- adjust to this node's GPU count + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + # Doc-level scoring (one doc per row, no packing): packed 512-token + # chunks contaminate per-doc gradients across doc boundaries and + # were found to suppress LDS (see gpt2_lotus_trackstar50q_docs). + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + approx_unrolling_cfg: + checkpoints: + - runs/lotus_source_q50/models/step_51 + - runs/lotus_source_q50/models/step_102 + - runs/lotus_source_q50/models/step_136 + - runs/lotus_source_q50/models/step_187 + - runs/lotus_source_q50/models/step_238 + - runs/lotus_source_q50/models/step_288 + segments: 3 + lr_list: [4.857423e-04, 5.616667e-04, 2.416667e-04] + step_size_list: [97, 96, 96] + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + # One score column per query (50 total), like the EK-FAC and + # TrackStar runs, for multi-query LDS validation. + query_aggregation: none + + - validate: + run_path: runs/lotus_source_q50_damp0_validate + model: gpt2 + overwrite: true + + scores: runs/lotus_source_q50_damp0/scores.npy + retrained_dir: runs/lotus + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 512 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + + batch_size: 64 diff --git a/examples/method_comparison/gpt2_wikitext_lotus_trackstar_q50.yaml b/examples/method_comparison/gpt2_wikitext_lotus_trackstar_q50.yaml new file mode 100644 index 00000000..e71f91b5 --- /dev/null +++ b/examples/method_comparison/gpt2_wikitext_lotus_trackstar_q50.yaml @@ -0,0 +1,170 @@ +# TrackStar (Chang et al., 2024) attribution scores for the lotus GPT-2 +# base model against the 50 test queries, validated against the complete +# 400-model lotus bank -- the same setting as the MAGIC (rho=0.92) and +# SOURCE (rho=0.39) comparisons. Doc-level scoring, projection_dim 32, +# Adam-normalized gradients from runs/lotus/optimizer.pt, standard +# query/value hessian mixing. + +# Run with `bergson examples/method_comparison/gpt2_wikitext_lotus_trackstar_q50.yaml` + +steps: +- trackstar: + index_cfg: + run_path: runs/lotus_trackstar_q50 + overwrite: true + model: runs/lotus/retrained/base + precision: fp32 + revision: null + distributed: + nnode: 1 + nproc_per_node: 8 + node_rank: null + fsdp: false + peft_init_kwargs: '' + model_kwargs: '' + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: train + subset: null + prompt_column: text + completion_column: '' + conversation_column: '' + reward_column: '' + skip_nan_rewards: false + truncation: false + format_template: '' + data_kwargs: '' + chunk_length: 0 + tokenizer: '' + drop_columns: true + max_tokens: null + use_tf32_matmuls: false + debug: false + projection_dim: 32 + include_bias: false + reshape_to_square: false + projection_type: rademacher + projection_target: per_module + token_batch_size: 1024 + max_batch_size: null + auto_batch_size: false + optimizer_state: runs/lotus/optimizer.pt + loss_fn: ce + loss_reduction: sum + label_smoothing: 0.0 + stream_shard_size: 400000 + split_attention_modules: [] + attention: + num_heads: 0 + head_size: 0 + head_dim: 0 + profile: false + filter_modules: lm_head + force_math_sdp: false + attribute_tokens: false + modules: [] + trackstar_cfg: + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: test[1:51] + subset: null + prompt_column: text + completion_column: '' + conversation_column: '' + reward_column: '' + skip_nan_rewards: false + truncation: false + format_template: '' + data_kwargs: '' + chunk_length: 0 + preprocess_cfg: + unit_normalize: true + hessian_path: null + inversion_cfg: + inversion: damped_inverse + damping_factor: 0.1 + aggregation: none + normalize_aggregated_grad: false + score_cfg: + query_path: '' + score: individual + batch_size: 1024 + precision: fp32 + modules: [] + higher_is_better: true + target_downweight_components: 1000 + stats_sample_size: 10000 + resume: true +- validate: + run_path: runs/lotus_trackstar_q50_validate + overwrite: true + model: gpt2 + precision: fp32 + revision: null + distributed: + nnode: 1 + nproc_per_node: 8 + node_rank: null + fsdp: false + peft_init_kwargs: '' + model_kwargs: '' + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: train + subset: null + prompt_column: text + completion_column: '' + conversation_column: '' + reward_column: '' + skip_nan_rewards: false + truncation: false + format_template: '' + data_kwargs: '' + chunk_length: 512 + tokenizer: '' + drop_columns: true + max_tokens: null + use_tf32_matmuls: false + debug: false + lr_schedule: + lr: 1.0e-05 + lr_scheduler_type: linear + lr_start: 0.0 + lr_end: 0.0 + warmup_steps: 0 + num_cycles: 0.5 + power: 1.0 + batch_size: 64 + num_epochs: 1 + seed: 42 + adam_beta1: 0.95 + adam_beta2: 0.975 + eps_root: 1.0e-08 + optimizer: adamw + save_optimizer_state: false + weight_decay: 0.01 + max_grad_norm: null + grad_checkpointing: false + resume: false + wandb_project: '' + save_retrained_models: false + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: test[1:51] + subset: null + prompt_column: text + completion_column: '' + conversation_column: '' + reward_column: '' + skip_nan_rewards: false + truncation: false + format_template: '' + data_kwargs: '' + chunk_length: 0 + query_method: mean + num_subsets: 100 + subset_strategy: random + exclude_zero_scores: false + subset_fraction: 0.0 + scores: runs/lotus_trackstar_q50/scores + retrained_dir: runs/lotus diff --git a/examples/pipelines/ekfac_qwen7b.yaml b/examples/pipelines/ekfac_qwen7b.yaml new file mode 100644 index 00000000..e2a970ae --- /dev/null +++ b/examples/pipelines/ekfac_qwen7b.yaml @@ -0,0 +1,59 @@ +# EKFAC influence pipeline for Qwen/Qwen2.5-7B-Instruct on a dummy test set. +# +# Fits an EK-FAC (KFAC + eigenvalue correction) Hessian over a small slice of +# pile-10k, then scores every index document against the mean gradient of a +# tiny 4-document query "test set". Small enough to sanity-check the pipeline +# end-to-end; scale up the dataset splits for a real run. +# +# Run with `bergson examples/pipelines/ekfac_qwen7b.yaml` +# +# Outputs land under runs/ekfac_qwen7b/. The score memmap is at +# runs/ekfac_qwen7b/scores/scores.bin; load it with bergson.data.load_scores. + +run_path: runs/ekfac_qwen7b +steps: + - ekfac: + index_cfg: + run_path: runs/ekfac_qwen7b + model: Qwen/Qwen2.5-7B-Instruct + overwrite: true + # bf16 keeps the 7B model comfortably within a single 48GB GPU. + precision: bf16 + token_batch_size: 2048 + # EK-FAC doesn't support gradient projection at any stage, so every + # gradient here is the full dense weight gradient. For a + # multi-billion-param model that means ONE document per batch — 16 + # docs would need hundreds of GB. Raise cautiously. + max_batch_size: 1 + # Qwen2.5's LM head is a huge (152k x 3584) nn.Linear; exclude it. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 + nnode: 1 + data: + # Dummy index corpus: a small slice of pile-10k. + dataset: NeelNanda/pile-10k + split: "train[:128]" + truncation: true + + hessian_cfg: + # KFAC + eigenvalue correction == EK-FAC. + method: kfac + ev_correction: true + + score_cfg: + batch_size: 512 + + preprocess_cfg: + unit_normalize: false + + hessian_pipeline_cfg: + query: + # Dummy query "test set": 8 held-out documents (>= nproc_per_node so + # every GPU is used; bergson caps workers to the dataset size). + dataset: NeelNanda/pile-10k + split: "train[128:136]" + truncation: true + query_aggregation: mean + inversion_cfg: + damping_factor: 0.1 diff --git a/examples/pipelines/recall_ekfac_olmo3.yaml b/examples/pipelines/recall_ekfac_olmo3.yaml new file mode 100644 index 00000000..ff9aac96 --- /dev/null +++ b/examples/pipelines/recall_ekfac_olmo3.yaml @@ -0,0 +1,63 @@ +# EK-FAC recall for the LoRA-finetuned Olmo-3-7B on the synthetic +# factual-recall data, to compare against an existing TrackStar result. +# +# EK-FAC (kfac + ev_correction) scores every statement against every question +# (query_aggregation: none -> one score column per question), then `recall` +# ranks the gold statements per question and reports MRR / Recall@10. +# +# Scale is 100 people (3600 statements x 400 questions) to match the TrackStar +# run already on disk at: +# runs/recall_trackstar_olmo3_100p_1gpu/eval/summary.csv +# so the two are directly comparable. (TrackStar there: MRR 0.547, Recall@10 +# 0.812.) The model (a public LoRA adapter over allenai/Olmo-3-7B-Instruct) and +# its base are public, so no HF token is needed. +# +# Prerequisites: +# 1. A node with working multi-GPU (NCCL healthy): EK-FAC shards its KFAC +# factors across GPUs and OOMs on a single 48GB card for a 7B model. No +# `distributed:` block is set, so all GPUs on the node are used. +# 2. The 100-person recall datasets (already present on our filesystem; if not, +# regenerate them deterministically): +# python -m bergson.recall.generate --num_people 100 --seed 0 --data_dir data +# +# Run with `python -m bergson examples/pipelines/recall_ekfac_olmo3.yaml`, then: +# cat runs/recall_ekfac_olmo3_100p/eval/summary.csv # EK-FAC +# cat runs/recall_trackstar_olmo3_100p_1gpu/eval/summary.csv # TrackStar +# +# For a larger comparison, regenerate both datasets at a new --num_people M, +# rerun TrackStar at that M, and replace every `100` / `_100p_` below with M. + +run_path: runs/recall_ekfac_olmo3_100p +steps: + - ekfac: + index_cfg: + run_path: runs/recall_ekfac_olmo3_100p/ekfac + model: EleutherAI/Olmo-3-7B-Instruct-recall-9000 + precision: bf16 + overwrite: true + force_math_sdp: true + data: + dataset: data/statements_100p_seed0.hf + prompt_column: fact + hessian_cfg: + method: kfac + ev_correction: true + score_cfg: + batch_size: 1024 + preprocess_cfg: {} + hessian_pipeline_cfg: + query: + dataset: data/questions_100p_seed0.hf + prompt_column: text + query_aggregation: none + inversion_cfg: + inversion: damped_inverse + damping_factor: 0.1 + + - recall: + run_path: runs/recall_ekfac_olmo3_100p/eval + scores: runs/recall_ekfac_olmo3_100p/ekfac/scores + k: 10 + data: + num_people: 100 + seed: 0 diff --git a/examples/pipelines/recall_ekfac_olmo3_fp32.yaml b/examples/pipelines/recall_ekfac_olmo3_fp32.yaml new file mode 100644 index 00000000..890df2fd --- /dev/null +++ b/examples/pipelines/recall_ekfac_olmo3_fp32.yaml @@ -0,0 +1,62 @@ +# EK-FAC recall for the LoRA-finetuned Olmo-3-7B on the synthetic +# factual-recall data, in fp32. +# +# Identical to recall_ekfac_olmo3.yaml except precision is fp32 (the model is +# loaded and gradients are collected in fp32 rather than bf16) and the run +# path has an _fp32 suffix. Compare against the fp32 TrackStar run at +# runs/recall_trackstar_olmo3_100p_fp32_1gpu/eval/summary.csv +# produced by recall_trackstar_olmo3_fp32.yaml. +# +# Prerequisites are the same as recall_ekfac_olmo3.yaml: a healthy multi-GPU +# node and the 100-person recall datasets: +# python -m bergson.recall.generate --num_people 100 --seed 0 --data_dir data +# +# Run with `python -m bergson examples/pipelines/recall_ekfac_olmo3_fp32.yaml`, +# then: +# cat runs/recall_ekfac_olmo3_100p_fp32/eval/summary.csv + +run_path: runs/recall_ekfac_olmo3_100p_fp32 +steps: + - ekfac: + index_cfg: + run_path: runs/recall_ekfac_olmo3_100p_fp32/ekfac + model: EleutherAI/Olmo-3-7B-Instruct-recall-9000 + precision: fp32 + overwrite: true + force_math_sdp: true + # Down from the 2048 default: in fp32 the 48GB A40s OOM during + # gradient collection at 2048, during the KFAC fit at 1024, and during + # scoring at 512 (per-sequence gradients for a batch are + # batch_size x 40M params fp32 next to the model + query slice). + token_batch_size: 256 + data: + dataset: data/statements_100p_seed0.hf + prompt_column: fact + hessian_cfg: + method: kfac + ev_correction: true + score_cfg: + batch_size: 1024 + # Score 64 queries per pass: the full 400-query set of unprojected + # fp32 gradients (~64GB) cannot sit on a 48GB card next to the model. + query_batch_size: 64 + preprocess_cfg: {} + hessian_pipeline_cfg: + # Skip pipeline steps whose output dir already exists (query gradients + # and KFAC factors survive a crash in a later step). + resume: true + query: + dataset: data/questions_100p_seed0.hf + prompt_column: text + query_aggregation: none + inversion_cfg: + inversion: damped_inverse + damping_factor: 0.1 + + - recall: + run_path: runs/recall_ekfac_olmo3_100p_fp32/eval + scores: runs/recall_ekfac_olmo3_100p_fp32/ekfac/scores + k: 10 + data: + num_people: 100 + seed: 0 diff --git a/examples/pipelines/recall_source_olmo3_fp32.yaml b/examples/pipelines/recall_source_olmo3_fp32.yaml new file mode 100644 index 00000000..49c94be1 --- /dev/null +++ b/examples/pipelines/recall_source_olmo3_fp32.yaml @@ -0,0 +1,79 @@ +# SOURCE (approximate unrolling, Bae et al.) recall for the LoRA-finetuned +# Olmo-3-7B on the synthetic factual-recall data, in fp32, to compare against +# the fp32 EK-FAC and TrackStar runs at the same 100-person scale: +# runs/recall_ekfac_olmo3_100p_fp32/eval/summary.csv +# runs/recall_trackstar_olmo3_100p_fp32_1gpu/eval/summary.csv +# +# 6 checkpoints / 3 segments over the 10125-step finetune +# (runs/finetune_recall_9000_olmo3: cosine LR peak 2e-4, 3% warmup, batch 128, +# 4 epochs). Checkpoints are HF adapter dirs exported from the trainer's DCP +# snapshots at steps {1700, 3400 | 5100, 6800 | 8400, 10124} — segment +# midpoints and ends for boundaries 0-3375-6750-10125: +# python scripts/export_dcp_checkpoints_to_hf.py \ +# --ckpt_dir runs/finetune_recall_9000_olmo3/checkpoints \ +# --base_model allenai/Olmo-3-7B-Instruct \ +# --peft_init_kwargs "r=16,lora_alpha=32,lora_dropout=0.0,target_modules=all-linear,task_type=CAUSAL_LM" \ +# --steps 1700 3400 5100 6800 8400 10124 \ +# --out_dir runs/recall_source_olmo3_100p_fp32/models +# +# lr_list is the per-third average of the training cosine schedule +# (LRScheduleConfig.get_schedule over 10125 steps); it cannot be inferred +# automatically because the MAGIC trainer writes no log_history.json. +# +# Disk: per-checkpoint covariances are ~66GB and segment eigenvectors ~66GB in +# fp32; expect a ~600GB peak. The per-ckpt covariance shards can be deleted +# once their segment's eigendecomposition (pipeline step 2/8) is done. +# +# Run with `python -m bergson examples/pipelines/recall_source_olmo3_fp32.yaml`, +# then: +# cat runs/recall_source_olmo3_100p_fp32/eval/summary.csv + +run_path: runs/recall_source_olmo3_100p_fp32 +steps: + - approxunrolling: + index_cfg: + run_path: runs/recall_source_olmo3_100p_fp32 + # Tokenizer + final-model reference; per-checkpoint passes load the + # adapter dirs below. + model: EleutherAI/Olmo-3-7B-Instruct-recall-9000 + # In the approxunrolling pipeline `overwrite` doubles as resume: + # steps whose output dirs exist are skipped after a crash. + overwrite: true + precision: fp32 + force_math_sdp: true + # fp32 memory ceiling on 48GB A40s, matching the EK-FAC fp32 run. + token_batch_size: 256 + data: + dataset: data/statements_100p_seed0.hf + prompt_column: fact + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + approx_unrolling_cfg: + checkpoints: + - runs/recall_source_olmo3_100p_fp32/models/step_1700 + - runs/recall_source_olmo3_100p_fp32/models/step_3400 + - runs/recall_source_olmo3_100p_fp32/models/step_5100 + - runs/recall_source_olmo3_100p_fp32/models/step_6800 + - runs/recall_source_olmo3_100p_fp32/models/step_8400 + - runs/recall_source_olmo3_100p_fp32/models/step_10124 + segments: 3 + lr_list: [1.770244e-04, 1.046429e-04, 1.833268e-05] + step_size_list: [3375, 3375, 3375] + query: + dataset: data/questions_100p_seed0.hf + prompt_column: text + # One score column per question, like the EK-FAC/TrackStar runs. + query_aggregation: none + # The 400 unaggregated fp32 query gradients (~60GB) cannot sit on a + # 48GB card next to the model during per-segment scoring. + query_batch_size: 64 + + - recall: + run_path: runs/recall_source_olmo3_100p_fp32/eval + scores: runs/recall_source_olmo3_100p_fp32/scores + k: 10 + data: + num_people: 100 + seed: 0 diff --git a/examples/pipelines/recall_trackstar_olmo3_fp32.yaml b/examples/pipelines/recall_trackstar_olmo3_fp32.yaml new file mode 100644 index 00000000..34821512 --- /dev/null +++ b/examples/pipelines/recall_trackstar_olmo3_fp32.yaml @@ -0,0 +1,57 @@ +# TrackStar recall for the LoRA-finetuned Olmo-3-7B on the synthetic +# factual-recall data, in fp32. +# +# Mirrors the config saved in runs/recall_trackstar_olmo3_100p_1gpu/config.yaml +# (bergson 0.10.0, git 989808d0) with precision fp32 instead of bf16 and an +# _fp32 run path. Kept single-process (nproc_per_node: 1) because the cosine +# recall variants OOM on host RAM when run multi-GPU. Compare against the fp32 +# EK-FAC run from recall_ekfac_olmo3_fp32.yaml. +# +# Run with `python -m bergson examples/pipelines/recall_trackstar_olmo3_fp32.yaml`, +# then: +# cat runs/recall_trackstar_olmo3_100p_fp32_1gpu/eval/summary.csv + +run_path: runs/recall_trackstar_olmo3_100p_fp32_1gpu +steps: + - trackstar: + index_cfg: + run_path: runs/recall_trackstar_olmo3_100p_fp32_1gpu/trackstar + overwrite: true + model: EleutherAI/Olmo-3-7B-Instruct-recall-9000 + precision: fp32 + distributed: + nnode: 1 + nproc_per_node: 1 + data: + dataset: data/statements_100p_seed0.hf + prompt_column: fact + projection_dim: 16 + projection_type: rademacher + projection_target: per_module + # Down from 2048 in the bf16 run: the fp32 model + backward OOMs a + # single 48GB A40 at that token budget. + token_batch_size: 512 + force_math_sdp: true + trackstar_cfg: + query: + dataset: data/questions_100p_seed0.hf + prompt_column: text + preprocess_cfg: + unit_normalize: true + inversion_cfg: + inversion: damped_inverse + damping_factor: 0.1 + aggregation: none + score_cfg: + score: individual + batch_size: 1024 + precision: fp32 + target_downweight_components: 1000 + + - recall: + run_path: runs/recall_trackstar_olmo3_100p_fp32_1gpu/eval + scores: runs/recall_trackstar_olmo3_100p_fp32_1gpu/trackstar/scores + k: 10 + data: + num_people: 100 + seed: 0 diff --git a/examples/pipelines/recall_trackstar_olmo3_fp32_dim32.yaml b/examples/pipelines/recall_trackstar_olmo3_fp32_dim32.yaml new file mode 100644 index 00000000..2cf661be --- /dev/null +++ b/examples/pipelines/recall_trackstar_olmo3_fp32_dim32.yaml @@ -0,0 +1,57 @@ +# TrackStar recall in fp32 with projection_dim 32 (vs 16 in +# recall_trackstar_olmo3_fp32.yaml, otherwise identical). Each module gradient +# gets a double-sided 32x32 random projection instead of 16x16, i.e. 4x the +# projected feature dims — testing whether TrackStar's gap to EK-FAC/SOURCE at +# dim 16 (see experiments/recall_attribution_fp32/README.md) is a projection +# bottleneck. dim 64 OOMs: the dense per-module Grams are (d^2)^2 fp32, ~30GB +# at d=64 next to the 28GB model. d=32 keeps them at ~1.9GB. +# +# Run with +# `python -m bergson examples/pipelines/recall_trackstar_olmo3_fp32_dim32.yaml`, +# then: +# cat runs/recall_trackstar_olmo3_100p_fp32_dim32/eval/summary.csv + +run_path: runs/recall_trackstar_olmo3_100p_fp32_dim32 +steps: + - trackstar: + index_cfg: + run_path: runs/recall_trackstar_olmo3_100p_fp32_dim32/trackstar + overwrite: true + model: EleutherAI/Olmo-3-7B-Instruct-recall-9000 + precision: fp32 + distributed: + nnode: 1 + nproc_per_node: 8 + data: + dataset: data/statements_100p_seed0.hf + prompt_column: fact + projection_dim: 32 + projection_type: rademacher + projection_target: per_module + # Down from 2048 in the bf16 run: the fp32 model + backward OOMs a + # single 48GB A40 at that token budget. + token_batch_size: 512 + force_math_sdp: true + trackstar_cfg: + query: + dataset: data/questions_100p_seed0.hf + prompt_column: text + preprocess_cfg: + unit_normalize: true + inversion_cfg: + inversion: damped_inverse + damping_factor: 0.1 + aggregation: none + score_cfg: + score: individual + batch_size: 1024 + precision: fp32 + target_downweight_components: 1000 + + - recall: + run_path: runs/recall_trackstar_olmo3_100p_fp32_dim32/eval + scores: runs/recall_trackstar_olmo3_100p_fp32_dim32/trackstar/scores + k: 10 + data: + num_people: 100 + seed: 0 diff --git a/tests/test_banked_model_loading.py b/tests/test_banked_model_loading.py index 5c4db953..0e2be00b 100644 --- a/tests/test_banked_model_loading.py +++ b/tests/test_banked_model_loading.py @@ -1,6 +1,6 @@ -"""Loading banked ``save_retrained_models`` checkpoints in ``evaluate_retrained``. +"""Loading banked ``save_models`` checkpoints in ``evaluate_retrained``. -``save_retrained_models`` banks whatever ``model.save_pretrained`` produces: +``save_models`` banks whatever ``model.save_pretrained`` produces: a full HF checkpoint for plain fine-tunes, or an adapter-only directory when ``peft_init_kwargs`` is set. ``AutoModelForCausalLM.from_pretrained`` pointed straight at an adapter-only directory silently returns a randomly-initialised From 40a8678f0ca4d4b405ffe8742898882ec19baf78 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 11:58:52 +0000 Subject: [PATCH 06/10] fixup --- .../approx_unrolling/adam_preconditioner.py | 50 ++++++++----------- 1 file changed, 20 insertions(+), 30 deletions(-) diff --git a/bergson/approx_unrolling/adam_preconditioner.py b/bergson/approx_unrolling/adam_preconditioner.py index 60edb276..c8a92f6c 100644 --- a/bergson/approx_unrolling/adam_preconditioner.py +++ b/bergson/approx_unrolling/adam_preconditioner.py @@ -1,15 +1,12 @@ -"""Per-segment Adam/AdamW diagonal preconditioners for the preconditioned -SOURCE (approx-unrolling) variant. +"""Per-segment Adam SOURCE preconditioners. :func:`build_segment_preconditioners` reads each checkpoint dir's -``optimizer.pt`` (written by ``TrainingConfig.save_optimizer_state``), -bias-corrects the second moments, averages them within each segment as the -K-FAC factors are, and writes +``optimizer.pt``, bias-corrects the second moments, averages them +within each segment, and writes P = 1 / (sqrt(v_hat_bar + eps_root) + eps) -to ``/segment_/preconditioner.safetensors``, oriented to bergson's -``[out, in]`` gradient grids. +to ``/segment_/preconditioner.safetensors``. """ from glob import glob @@ -31,11 +28,11 @@ def _module_grids(hessian_dir: str | Path) -> dict[str, tuple[int, int]]: - """Module name -> full [out, in] grid shape, read from the factor shards. + """Module name -> [out, in] for the module's weight matrix shape. - The activation eigenvectors are ``[c_i, I]`` (columns = full in-dim) and - the gradient eigenvectors ``[c_o, O]``, so shard_0's column counts give - the full grid shape regardless of world size. + The activation eigenvectors are ``[c_i, I]`` and the gradient + eigenvectors are ``[c_o, O]``, so any shard's column dimensions + give the full weight matrix shape. """ grids: dict[str, tuple[int, int]] = {} a_dir = Path(hessian_dir) / "eigen_activation_sharded" @@ -56,8 +53,7 @@ def _module_grids(hessian_dir: str | Path) -> dict[str, tuple[int, int]]: def _match_params( grids: dict[str, tuple[int, int]], param_names: list[str] ) -> dict[str, str]: - """Factor module name -> param name. Factor names are suffixes of the - model's param names (e.g. ``h.0.attn.c_attn`` -> + """Factor module name -> param name, e.g. ``h.0.attn.c_attn`` -> ``transformer.h.0.attn.c_attn.weight``).""" matches: dict[str, str] = {} for module in grids: @@ -73,8 +69,7 @@ def _match_params( def _orient(p: torch.Tensor, grid: tuple[int, int], module: str, layer) -> torch.Tensor: - """Orient a param-shaped grid to bergson's [out, in] layout, raising on a - shape that fits neither orientation.""" + """Orient a param-shaped grid to bergson's [out, in] layout.""" oriented = orient_second_moment(p, *grid, layer=layer) if oriented is None: raise ValueError( @@ -85,20 +80,17 @@ def _orient(p: torch.Tensor, grid: tuple[int, int], module: str, layer) -> torch def _adam_hparams(optimizer_pt: dict) -> tuple[float, float, float]: - """(beta2, eps, eps_root) from a standard ``optimizer.pt``'s param_groups. + """Get (beta2, eps, eps_root) from a standard ``optimizer.pt``. ``betas`` and ``eps`` are standard AdamW param-group keys (HF Trainer - checkpoints carry them); ``eps_root`` is torchopt-specific and defaults - to 0 (standard AdamW has no epsilon inside the sqrt).""" + saves them); ``eps_root`` is torchopt-specific and defaults to 0.""" groups = optimizer_pt.get("param_groups", []) betas = groups[0].get("betas") if groups else None eps = groups[0].get("eps") if groups else None if betas is None or eps is None: raise ValueError( - "optimizer.pt has no param_groups betas/eps; regenerate the " - "snapshot exports with TrainingConfig.save_optimizer_state " - "(older exports lack the hyperparameters needed to build the " - "preconditioner)." + "optimizer.pt has no param_groups betas/eps; save " + "with TrainingConfig.save_optimizer_state " ) return float(betas[1]), float(eps), float(groups[0].get("eps_root", 0.0)) @@ -106,8 +98,8 @@ def _adam_hparams(optimizer_pt: dict) -> tuple[float, float, float]: def _extract_v_hat( optimizer_pt: dict, index_to_name: dict[int, str], beta2: float ) -> dict[str, torch.Tensor]: - """Bias-corrected second moments keyed by param name from a standard - ``optimizer.pt`` dict carrying a per-entry ``step``.""" + """Bias-corrected second moments keyed by param name from an + ``optimizer.pt`` with a per-entry ``step``.""" v_hat: dict[str, torch.Tensor] = {} for idx, entry in optimizer_pt["state"].items(): @@ -119,8 +111,8 @@ def _extract_v_hat( continue if "step" not in entry: raise ValueError( - f"optimizer.pt state entry {idx} has no step; regenerate the " - "snapshot exports with TrainingConfig.save_optimizer_state." + f"optimizer.pt state entry {idx} has no step; save " + "with TrainingConfig.save_optimizer_state." ) step = int( entry["step"].item() if torch.is_tensor(entry["step"]) else entry["step"] @@ -137,10 +129,8 @@ def build_segment_preconditioners( segments: int, ) -> list[Path]: """Build ``/segment_/preconditioner.safetensors`` for every - segment from the checkpoints' ``optimizer.pt`` files; see module docs. - The Adam hyperparameters (betas/eps/eps_root) are read from the files. - - Returns the per-segment preconditioner paths. + segment from the checkpoints' ``optimizer.pt``. Returns the + per-segment preconditioner paths. """ base = Path(run_path) per_segment = len(checkpoints) // segments From 529232bda8fac2e83884cd88a9533220eb6a6bac Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 12:05:48 +0000 Subject: [PATCH 07/10] feat: utility to place optimizer states where SOURCE and TrackStar read them The trainer writes each snapshot's second moments as a sibling file, /step_.optimizer.pt, next to that step's step_.ckpt. Both consumers instead want optimizer.pt *inside* a directory: SOURCE reads /optimizer.pt under use_adam_preconditioner, and TrackStar resolves PreprocessConfig.optimizer_state through load_optimizer, which loads optimizer.pt from a directory. Exporting to HF-format checkpoint dirs does not carry the sibling files across, so both paths previously told the user to copy the files by hand. place_optimizer_states() does that placement for a run's checkpoints, and place_final_optimizer_state() does it for the end-of-training state that TrackStar normalizes with. Since both consumers agree on /optimizer.pt, one call satisfies either. States are matched to destinations by step index when the destination names one, else positionally in step order with the counts required to agree. A half-named list, an unknown step, or a count mismatch raises: silently mismatching would build a preconditioner from the wrong step's moments, which no downstream check would catch. Existing files are never clobbered without overwrite=True, and the check precedes all writes so a refusal leaves nothing half-placed. Defaults to relative symlinks -- second moments are one scalar per weight, so copying one per checkpoint is wasteful, and relative keeps an exported tree valid when moved. copy/hardlink/move are available for self-contained exports. OPTIMIZER_STATE_FILE now has a single definition, re-exported from adam_preconditioner so pipeline.py's import is unchanged. --- .../approx_unrolling/adam_preconditioner.py | 14 +- .../approx_unrolling/approx_unrolling_math.py | 10 +- bergson/approx_unrolling/pipeline.py | 5 +- bergson/utils/optimizer_placement.py | 204 +++++++++++++++++ tests/test_optimizer_placement.py | 205 ++++++++++++++++++ 5 files changed, 425 insertions(+), 13 deletions(-) create mode 100644 bergson/utils/optimizer_placement.py create mode 100644 tests/test_optimizer_placement.py diff --git a/bergson/approx_unrolling/adam_preconditioner.py b/bergson/approx_unrolling/adam_preconditioner.py index c8a92f6c..c9b0ba1f 100644 --- a/bergson/approx_unrolling/adam_preconditioner.py +++ b/bergson/approx_unrolling/adam_preconditioner.py @@ -1,7 +1,7 @@ """Per-segment Adam SOURCE preconditioners. :func:`build_segment_preconditioners` reads each checkpoint dir's -``optimizer.pt``, bias-corrects the second moments, averages them +``optimizer.pt``, bias-corrects the second moments, averages them within each segment, and writes P = 1 / (sqrt(v_hat_bar + eps_root) + eps) @@ -23,14 +23,15 @@ optimizer_param_index_to_name, orient_second_moment, ) +from ..utils.optimizer_placement import OPTIMIZER_STATE_FILE -OPTIMIZER_STATE_FILE = "optimizer.pt" +__all__ = ["OPTIMIZER_STATE_FILE", "build_segment_preconditioners"] def _module_grids(hessian_dir: str | Path) -> dict[str, tuple[int, int]]: """Module name -> [out, in] for the module's weight matrix shape. - The activation eigenvectors are ``[c_i, I]`` and the gradient + The activation eigenvectors are ``[c_i, I]`` and the gradient eigenvectors are ``[c_o, O]``, so any shard's column dimensions give the full weight matrix shape. """ @@ -129,7 +130,7 @@ def build_segment_preconditioners( segments: int, ) -> list[Path]: """Build ``/segment_/preconditioner.safetensors`` for every - segment from the checkpoints' ``optimizer.pt``. Returns the + segment from the checkpoints' ``optimizer.pt``. Returns the per-segment preconditioner paths. """ base = Path(run_path) @@ -148,8 +149,9 @@ def build_segment_preconditioners( raise FileNotFoundError( f"use_adam_preconditioner requires " f"{Path(ckpt) / OPTIMIZER_STATE_FILE}; write it during " - "training with TrainingConfig.save_optimizer_state and " - "copy it into the exported checkpoint dir." + "training with TrainingConfig.save_optimizer_state, then " + "place it with bergson.utils.optimizer_placement." + "place_optimizer_states(save_dir, checkpoints)." ) optimizer_pt = load_optimizer(str(ckpt)) beta2, eps, eps_root = _adam_hparams(optimizer_pt) diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index b607392d..648f88c5 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -152,11 +152,11 @@ def apply_eigfn_to_query( ``fn_kind`` is "f_segment" or "f_backward". The segment eigenvalues are already checkpoint-averaged (expected eigenvalues), so the eigenfunction is - applied to them directly. ``preconditioner_path`` selects the - preconditioned-optimizer variant: the eigenfunction is evaluated on a - diagonal approximation of P^1/2 H P^1/2 in parameter space, with - F_segment's output additionally multiplied by P (its P^1/2 . P^1/2 - sandwich; F_backward's sandwich cancels).""" + applied to them directly. + ``preconditioner_path`` selects the preconditioned-optimizer variant: the + eigenfunction is evaluated on a diagonal approximation of P^1/2 H P^1/2 in + parameter space, with F_segment's output additionally multiplied by P + (its P^1/2 . P^1/2 sandwich; F_backward's sandwich cancels).""" cfg = EkfacConfig( hessian_method_path=str(segment_dir), gradient_path=str(src_grad_path), diff --git a/bergson/approx_unrolling/pipeline.py b/bergson/approx_unrolling/pipeline.py index dbc9902d..99f1b2d5 100644 --- a/bergson/approx_unrolling/pipeline.py +++ b/bergson/approx_unrolling/pipeline.py @@ -115,8 +115,9 @@ def approx_unrolling_pipeline( if not state_path.exists(): raise FileNotFoundError( f"use_adam_preconditioner requires {state_path}; write it " - "during training with TrainingConfig.save_optimizer_state " - "and copy it into the exported checkpoint dir." + "during training with TrainingConfig.save_optimizer_state, " + "then place it with bergson.utils.optimizer_placement." + "place_optimizer_states(save_dir, checkpoints)." ) lr_times_steps_per_segment = compute_lr_times_steps_per_segment( diff --git a/bergson/utils/optimizer_placement.py b/bergson/utils/optimizer_placement.py new file mode 100644 index 00000000..17f202b8 --- /dev/null +++ b/bergson/utils/optimizer_placement.py @@ -0,0 +1,204 @@ +"""Place trainer-written optimizer states where attribution methods look for them. + +The trainer writes each snapshot's second moments as a *sibling* file, +``/step_.optimizer.pt``, next to that step's ``step_.ckpt`` +(see :meth:`bergson.magic.trainer.Trainer.train`). Both attribution consumers +instead expect ``optimizer.pt`` *inside* a directory: + +- SOURCE / approximate unrolling reads ``/optimizer.pt`` for every + checkpoint when ``ApproxUnrollingConfig.use_adam_preconditioner`` is set. +- TrackStar (and any ``build``) resolves ``PreprocessConfig.optimizer_state`` + via :func:`bergson.utils.load_from_optimizer.load_optimizer`, which accepts a + directory and loads ``optimizer.pt`` from inside it. + +Exporting a run to HF-format checkpoint dirs does not carry the sibling files +across, which is what previously forced a manual copy. :func:`place_optimizer_states` +does that placement, and satisfies both consumers with one call since they agree +on the ``/optimizer.pt`` layout. +""" + +import os +import re +import shutil +from pathlib import Path +from typing import Literal, Sequence + +OPTIMIZER_STATE_FILE = "optimizer.pt" +"""Filename both SOURCE and TrackStar look for inside a checkpoint directory.""" + +_STEP_OPTIMIZER_RE = re.compile(r"^step_(\d+)\.optimizer\.pt$") +_STEP_DIR_RE = re.compile(r"step_(\d+)") + +PlacementMode = Literal["copy", "symlink", "hardlink", "move"] + + +def sorted_optimizer_states(save_dir: str | Path) -> list[tuple[int, Path]]: + """Return ``(step, path)`` for each ``step_.optimizer.pt``, sorted by step. + + Mirrors :func:`bergson.data.sorted_checkpoints`, which does the same for the + ``step_.ckpt`` directories these files sit beside. + """ + save_dir = Path(save_dir) + if not save_dir.is_dir(): + raise NotADirectoryError(f"{save_dir} is not a directory") + + found = [] + for name in os.listdir(save_dir): + match = _STEP_OPTIMIZER_RE.match(name) + if match: + found.append((int(match.group(1)), save_dir / name)) + + return sorted(found, key=lambda pair: pair[0]) + + +def _step_of(path: Path) -> int | None: + """Step index named by a checkpoint directory, or None if it names none.""" + match = _STEP_DIR_RE.search(path.name) + return int(match.group(1)) if match else None + + +def _place(src: Path, dst: Path, mode: PlacementMode) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + if dst.exists() or dst.is_symlink(): + dst.unlink() + + match mode: + case "copy": + shutil.copy2(src, dst) + case "move": + shutil.move(str(src), str(dst)) + case "symlink": + # Relative, so an exported run stays valid if the tree is moved. + dst.symlink_to(os.path.relpath(src, dst.parent)) + case "hardlink": + os.link(src, dst) + case other: + raise ValueError(f"Unsupported placement mode: {other}") + + +def place_optimizer_states( + save_dir: str | Path, + checkpoints: Sequence[str | Path], + *, + mode: PlacementMode = "symlink", + overwrite: bool = False, +) -> dict[Path, Path]: + """Put an ``optimizer.pt`` inside each checkpoint dir, from ``save_dir``. + + ``save_dir`` is the trainer's output directory, holding the + ``step_.optimizer.pt`` files written when + ``TrainingConfig.save_optimizer_state`` is set. ``checkpoints`` are the + destination directories -- typically the exported HF checkpoints listed in + ``ApproxUnrollingConfig.checkpoints``. + + States are matched to destinations by step index when the destination + directory names one (``step_12`` / ``step_12.ckpt``); otherwise they are + matched positionally in step order, which requires the counts to be equal. + Mixing the two is an error, since a partial name match usually means the + caller passed a directory list that does not correspond to this run. + + ``mode`` defaults to ``"symlink"``: these files are large (second moments + are one scalar per weight) and duplicating them per checkpoint is wasteful. + Use ``"copy"`` when the destinations must be self-contained, e.g. before + uploading them somewhere. + + Returns ``{destination optimizer.pt: source file}``. Raises rather than + silently skipping, so a wrong directory list fails loudly instead of + producing a preconditioner from the wrong steps. + """ + states = sorted_optimizer_states(save_dir) + if not states: + raise FileNotFoundError( + f"No step_.optimizer.pt files in {save_dir}. Train with " + "TrainingConfig.save_optimizer_state=True to write them." + ) + + dsts = [Path(c) for c in checkpoints] + if not dsts: + raise ValueError("checkpoints is empty; nothing to place") + + missing = [d for d in dsts if not d.is_dir()] + if missing: + raise NotADirectoryError( + f"Checkpoint directories do not exist: {[str(m) for m in missing]}" + ) + + named = [_step_of(d) for d in dsts] + if all(step is not None for step in named): + steps = [step for step in named if step is not None] + by_step = dict(states) + unknown = [step for step in steps if step not in by_step] + if unknown: + raise FileNotFoundError( + f"No step_.optimizer.pt in {save_dir} for steps {unknown}; " + f"available steps: {sorted(by_step)}" + ) + pairs = [(by_step[step], dst) for step, dst in zip(steps, dsts)] + elif any(step is not None for step in named): + raise ValueError( + "Some checkpoint directories name a step and some do not, so they " + "cannot be matched reliably; pass directories that all name their " + "step, or none that do (for positional matching)." + ) + else: + if len(states) != len(dsts): + raise ValueError( + f"Cannot match positionally: {len(states)} optimizer states in " + f"{save_dir} but {len(dsts)} checkpoint directories. Name the " + "destination dirs after their steps to match by step instead." + ) + pairs = [(src, dst) for (_, src), dst in zip(states, dsts)] + + if not overwrite: + existing = [dst / OPTIMIZER_STATE_FILE for _, dst in pairs] + clashes = [p for p in existing if p.exists() or p.is_symlink()] + if clashes: + raise FileExistsError( + f"{[str(c) for c in clashes]} already exist; pass overwrite=True " + "to replace them." + ) + + placed: dict[Path, Path] = {} + for src, dst in pairs: + target = dst / OPTIMIZER_STATE_FILE + _place(src, target, mode) + placed[target] = src + + return placed + + +def place_final_optimizer_state( + run_path: str | Path, + destination: str | Path, + *, + mode: PlacementMode = "symlink", + overwrite: bool = False, +) -> Path: + """Put the run's final ``optimizer.pt`` inside ``destination``. + + ``TrainingConfig.save_optimizer_state`` writes the end-of-training state to + ``/optimizer.pt``. This places it inside an exported model + directory so ``PreprocessConfig.optimizer_state`` can point at that + directory -- the TrackStar gradient-normalization path. + + Returns the path written. + """ + src = Path(run_path) / OPTIMIZER_STATE_FILE + if not src.is_file(): + raise FileNotFoundError( + f"{src} not found. Train with TrainingConfig.save_optimizer_state=True " + "to write it." + ) + + dst_dir = Path(destination) + if not dst_dir.is_dir(): + raise NotADirectoryError(f"{dst_dir} is not a directory") + + target = dst_dir / OPTIMIZER_STATE_FILE + if target.resolve() == src.resolve(): + return target + if not overwrite and (target.exists() or target.is_symlink()): + raise FileExistsError(f"{target} already exists; pass overwrite=True.") + + _place(src, target, mode) + return target diff --git a/tests/test_optimizer_placement.py b/tests/test_optimizer_placement.py new file mode 100644 index 00000000..373c7186 --- /dev/null +++ b/tests/test_optimizer_placement.py @@ -0,0 +1,205 @@ +"""Placing trainer-written optimizer states where SOURCE and TrackStar read them. + +The trainer writes ``step_.optimizer.pt`` as a sibling of ``step_.ckpt``; +both consumers want ``optimizer.pt`` *inside* a directory. These tests pin the +matching rules, since a wrong match silently builds a preconditioner from the +wrong step rather than failing. +""" + +import os +import shutil +from pathlib import Path + +import pytest +import torch + +from bergson.approx_unrolling.adam_preconditioner import OPTIMIZER_STATE_FILE +from bergson.utils.load_from_optimizer import load_optimizer +from bergson.utils.optimizer_placement import ( + place_final_optimizer_state, + place_optimizer_states, + sorted_optimizer_states, +) + + +def _write_state(path: Path, step: int) -> None: + """A minimal optimizer.pt whose payload identifies its step.""" + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "state": {0: {"exp_avg_sq": torch.full((2, 3), float(step))}}, + "param_groups": [{"betas": (0.9, 0.999), "eps": 1e-8, "lr": 1e-4}], + }, + path, + ) + + +@pytest.fixture +def save_dir(tmp_path) -> Path: + """A trainer output dir with sibling states at steps 0, 4 and 9.""" + d = tmp_path / "run" + for step in (0, 4, 9): + _write_state(d / f"step_{step}.optimizer.pt", step) + (d / f"step_{step}.ckpt").mkdir(parents=True, exist_ok=True) + return d + + +def _exported(root: Path, names) -> list[Path]: + dirs = [] + for name in names: + p = root / name + p.mkdir(parents=True, exist_ok=True) + dirs.append(p) + return dirs + + +def test_sorted_optimizer_states_orders_numerically(save_dir): + """Steps must sort numerically, not lexically (step_10 < step_9 as strings).""" + _write_state(save_dir / "step_10.optimizer.pt", 10) + assert [s for s, _ in sorted_optimizer_states(save_dir)] == [0, 4, 9, 10] + + +def test_sorted_optimizer_states_ignores_other_files(save_dir): + """The final optimizer.pt and the .ckpt dirs are not per-step states.""" + _write_state(save_dir / OPTIMIZER_STATE_FILE, -1) + assert [s for s, _ in sorted_optimizer_states(save_dir)] == [0, 4, 9] + + +def test_places_by_step_name(save_dir, tmp_path): + """Directories naming a step get that step's state, regardless of order.""" + dsts = _exported(tmp_path / "hf", ["step_9", "step_0", "step_4"]) + placed = place_optimizer_states(save_dir, dsts, mode="copy") + + assert len(placed) == 3 + for dst, expected in zip(dsts, (9, 0, 4)): + blob = load_optimizer(str(dst)) + assert blob["state"][0]["exp_avg_sq"][0, 0].item() == expected + + +def test_places_positionally_when_unnamed(save_dir, tmp_path): + """Unnamed dirs are matched in step order, so seg 0 gets the earliest step.""" + dsts = _exported(tmp_path / "hf", ["a", "b", "c"]) + place_optimizer_states(save_dir, dsts, mode="copy") + + for dst, expected in zip(dsts, (0, 4, 9)): + blob = load_optimizer(str(dst)) + assert blob["state"][0]["exp_avg_sq"][0, 0].item() == expected + + +def test_positional_requires_equal_counts(save_dir, tmp_path): + """A short/long dir list means the caller passed the wrong run; don't guess.""" + dsts = _exported(tmp_path / "hf", ["a", "b"]) + with pytest.raises(ValueError, match="Cannot match positionally"): + place_optimizer_states(save_dir, dsts, mode="copy") + + +def test_mixed_naming_is_rejected(save_dir, tmp_path): + """Half-named lists can't be matched reliably, so refuse rather than guess.""" + dsts = _exported(tmp_path / "hf", ["step_0", "b", "c"]) + with pytest.raises(ValueError, match="name a step and some do not"): + place_optimizer_states(save_dir, dsts, mode="copy") + + +def test_unknown_step_is_rejected(save_dir, tmp_path): + """A named step with no matching state is an error, not a skip.""" + dsts = _exported(tmp_path / "hf", ["step_0", "step_7"]) + with pytest.raises(FileNotFoundError, match=r"steps \[7\]"): + place_optimizer_states(save_dir, dsts, mode="copy") + + +def test_refuses_to_clobber_without_overwrite(save_dir, tmp_path): + """Existing optimizer.pt files are protected; nothing is written on refusal.""" + dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) + (dsts[1] / OPTIMIZER_STATE_FILE).write_text("do not clobber") + + with pytest.raises(FileExistsError): + place_optimizer_states(save_dir, dsts, mode="copy") + + assert (dsts[1] / OPTIMIZER_STATE_FILE).read_text() == "do not clobber" + # The check precedes all writes, so the untouched dirs stay untouched. + assert not (dsts[0] / OPTIMIZER_STATE_FILE).exists() + + place_optimizer_states(save_dir, dsts, mode="copy", overwrite=True) + assert load_optimizer(str(dsts[1]))["state"][0]["exp_avg_sq"][0, 0].item() == 4 + + +def test_symlink_mode_is_relative_and_loadable(save_dir, tmp_path): + """Default mode links rather than duplicating a state per checkpoint.""" + dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) + place_optimizer_states(save_dir, dsts) + + link = dsts[0] / OPTIMIZER_STATE_FILE + assert link.is_symlink() + assert not os.path.isabs(os.readlink(link)), "link must be relative" + assert load_optimizer(str(dsts[0]))["state"][0]["exp_avg_sq"][0, 0].item() == 0 + + +def test_symlinks_survive_relocation(save_dir, tmp_path): + """Relative links keep resolving when run and checkpoints move together.""" + dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) + place_optimizer_states(save_dir, dsts) + + moved = tmp_path / "moved" + moved.mkdir() + shutil.copytree(tmp_path / "hf", moved / "hf", symlinks=True) + shutil.copytree(save_dir, moved / save_dir.name, symlinks=True) + + relocated = moved / "hf" / "step_0" / OPTIMIZER_STATE_FILE + assert relocated.is_symlink() + assert relocated.resolve().is_file(), "relative link broke after relocation" + assert load_optimizer(str(moved / "hf" / "step_0"))["state"][0][ + "exp_avg_sq" + ][0, 0].item() == 0 + + +def test_hardlink_and_move_modes(save_dir, tmp_path): + dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) + place_optimizer_states(save_dir, dsts, mode="hardlink") + assert load_optimizer(str(dsts[2]))["state"][0]["exp_avg_sq"][0, 0].item() == 9 + + dsts2 = _exported(tmp_path / "hf2", ["step_0", "step_4", "step_9"]) + place_optimizer_states(save_dir, dsts2, mode="move") + assert not (save_dir / "step_9.optimizer.pt").exists() + assert load_optimizer(str(dsts2[2]))["state"][0]["exp_avg_sq"][0, 0].item() == 9 + + +def test_missing_states_names_the_config_flag(tmp_path): + """The error should say how to produce the files, not just that they're absent.""" + empty = tmp_path / "empty" + empty.mkdir() + dsts = _exported(tmp_path / "hf", ["step_0"]) + with pytest.raises(FileNotFoundError, match="save_optimizer_state"): + place_optimizer_states(empty, dsts) + + +def test_missing_destination_is_rejected(save_dir, tmp_path): + with pytest.raises(NotADirectoryError): + place_optimizer_states(save_dir, [tmp_path / "nope"]) + + +def test_place_final_state_for_trackstar(save_dir, tmp_path): + """TrackStar points optimizer_state at a dir; the final state must land there.""" + _write_state(save_dir / OPTIMIZER_STATE_FILE, 99) + exported = _exported(tmp_path / "model", ["final"])[0] + + target = place_final_optimizer_state(save_dir, exported, mode="copy") + + assert target == exported / OPTIMIZER_STATE_FILE + # load_optimizer resolves a directory, which is what PreprocessConfig takes. + assert load_optimizer(str(exported))["state"][0]["exp_avg_sq"][0, 0].item() == 99 + + +def test_place_final_state_is_idempotent_in_place(save_dir): + """Pointing it at its own directory is a no-op, not a self-clobber.""" + _write_state(save_dir / OPTIMIZER_STATE_FILE, 99) + target = place_final_optimizer_state(save_dir, save_dir) + assert target.is_file() + assert load_optimizer(str(save_dir))["state"][0]["exp_avg_sq"][0, 0].item() == 99 + + +def test_place_final_state_missing_names_the_flag(tmp_path): + empty = tmp_path / "empty" + empty.mkdir() + dst = _exported(tmp_path / "model", ["final"])[0] + with pytest.raises(FileNotFoundError, match="save_optimizer_state"): + place_final_optimizer_state(empty, dst) From 298f2621ce23434aed4466572e5b4e90725f695b Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 14:49:40 +0000 Subject: [PATCH 08/10] chore: drop files that do not belong on this branch The examples/ additions were untracked scratch swept in by a careless git add -A during the rebase. optimizer_placement.py and its tests were superseded on main, which keeps each checkpoint's optimizer state inside the checkpoint rather than matching loose files back to steps. --- .../approx_unrolling/adam_preconditioner.py | 7 +- .../approx_unrolling/approx_unrolling_math.py | 10 - bergson/approx_unrolling/pipeline.py | 5 +- bergson/magic/cli.py | 12 - bergson/utils/optimizer_placement.py | 204 ----------------- examples/bank_baselines/README.md | 41 ---- examples/bank_baselines/__init__.py | 0 .../bank_baselines/activation_baseline.py | 153 ------------- examples/bank_baselines/bm25_baseline.py | 69 ------ examples/bank_baselines/common.py | 156 ------------- examples/bank_baselines/gradient_baseline.py | 92 -------- examples/bank_baselines/qwen3_baseline.py | 76 ------- examples/bank_baselines/semantic_baseline.py | 102 --------- examples/magic/gpt2_wikitext_bank.yaml | 43 ---- examples/magic/gpt2_wikitext_dropout.yaml | 40 ---- examples/magic/gpt2_wikitext_dropout_s15.yaml | 42 ---- examples/magic/gpt2_wikitext_lotus.yaml | 44 ---- .../magic/gpt2_wikitext_lotus_ekfac_q50.yaml | 79 ------- .../gpt2_wikitext_lotus_final_q01_50.yaml | 31 --- examples/magic/gpt2_wikitext_lotus_mq.yaml | 70 ------ .../gpt2_wikitext_train_epsroot0_bs64.yaml | 54 ----- .../gpt2_wikitext_train_stdadam_bs64.yaml | 54 ----- .../gpt2_wikitext_lotus_source_adam_q50.yaml | 91 -------- .../gpt2_wikitext_lotus_source_q50.yaml | 79 ------- .../gpt2_wikitext_lotus_trackstar_q50.yaml | 170 --------------- examples/pipelines/ekfac_qwen7b.yaml | 59 ----- examples/pipelines/recall_ekfac_olmo3.yaml | 63 ------ .../pipelines/recall_ekfac_olmo3_fp32.yaml | 62 ------ .../pipelines/recall_source_olmo3_fp32.yaml | 79 ------- .../recall_trackstar_olmo3_fp32.yaml | 57 ----- .../recall_trackstar_olmo3_fp32_dim32.yaml | 57 ----- tests/test_optimizer_placement.py | 205 ------------------ 32 files changed, 5 insertions(+), 2301 deletions(-) delete mode 100644 bergson/utils/optimizer_placement.py delete mode 100644 examples/bank_baselines/README.md delete mode 100644 examples/bank_baselines/__init__.py delete mode 100644 examples/bank_baselines/activation_baseline.py delete mode 100644 examples/bank_baselines/bm25_baseline.py delete mode 100644 examples/bank_baselines/common.py delete mode 100644 examples/bank_baselines/gradient_baseline.py delete mode 100644 examples/bank_baselines/qwen3_baseline.py delete mode 100644 examples/bank_baselines/semantic_baseline.py delete mode 100644 examples/magic/gpt2_wikitext_bank.yaml delete mode 100644 examples/magic/gpt2_wikitext_dropout.yaml delete mode 100644 examples/magic/gpt2_wikitext_dropout_s15.yaml delete mode 100644 examples/magic/gpt2_wikitext_lotus.yaml delete mode 100644 examples/magic/gpt2_wikitext_lotus_ekfac_q50.yaml delete mode 100644 examples/magic/gpt2_wikitext_lotus_final_q01_50.yaml delete mode 100644 examples/magic/gpt2_wikitext_lotus_mq.yaml delete mode 100644 examples/magic/gpt2_wikitext_train_epsroot0_bs64.yaml delete mode 100644 examples/magic/gpt2_wikitext_train_stdadam_bs64.yaml delete mode 100644 examples/method_comparison/gpt2_wikitext_lotus_source_adam_q50.yaml delete mode 100644 examples/method_comparison/gpt2_wikitext_lotus_source_q50.yaml delete mode 100644 examples/method_comparison/gpt2_wikitext_lotus_trackstar_q50.yaml delete mode 100644 examples/pipelines/ekfac_qwen7b.yaml delete mode 100644 examples/pipelines/recall_ekfac_olmo3.yaml delete mode 100644 examples/pipelines/recall_ekfac_olmo3_fp32.yaml delete mode 100644 examples/pipelines/recall_source_olmo3_fp32.yaml delete mode 100644 examples/pipelines/recall_trackstar_olmo3_fp32.yaml delete mode 100644 examples/pipelines/recall_trackstar_olmo3_fp32_dim32.yaml delete mode 100644 tests/test_optimizer_placement.py diff --git a/bergson/approx_unrolling/adam_preconditioner.py b/bergson/approx_unrolling/adam_preconditioner.py index c9b0ba1f..05e9902e 100644 --- a/bergson/approx_unrolling/adam_preconditioner.py +++ b/bergson/approx_unrolling/adam_preconditioner.py @@ -23,7 +23,7 @@ optimizer_param_index_to_name, orient_second_moment, ) -from ..utils.optimizer_placement import OPTIMIZER_STATE_FILE +from ..utils.trainer_export import OPTIMIZER_STATE_FILE __all__ = ["OPTIMIZER_STATE_FILE", "build_segment_preconditioners"] @@ -149,9 +149,8 @@ def build_segment_preconditioners( raise FileNotFoundError( f"use_adam_preconditioner requires " f"{Path(ckpt) / OPTIMIZER_STATE_FILE}; write it during " - "training with TrainingConfig.save_optimizer_state, then " - "place it with bergson.utils.optimizer_placement." - "place_optimizer_states(save_dir, checkpoints)." + "training with save_optimizer_state='all', then export " + "with bergson.utils.trainer_export.export_checkpoints." ) optimizer_pt = load_optimizer(str(ckpt)) beta2, eps, eps_root = _adam_hparams(optimizer_pt) diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index 648f88c5..4ec7b518 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -60,7 +60,6 @@ def compute_lr_times_steps_per_segment( ) -> list[float]: """Per-segment lr * K. Use ``lr_list * step_size_list`` if set on config; else equal-partition log_history.json into segments and sum per-step LRs. -<<<<<<< HEAD With SGD heavy-ball ``momentum`` beta, scale by the terminal velocity 1/(1-beta) (Bae et al. 2024, App. D.2). @@ -72,15 +71,6 @@ def compute_lr_times_steps_per_segment( raise ValueError(f"momentum must be in [0, 1), got {momentum}.") momentum_scale = 1.0 / (1.0 - momentum) -======= - With SGD heavy-ball ``momentum`` beta, scale by the terminal velocity - 1/(1-beta) (Bae et al., 2024, Appendix D.2).""" - cfg = approx_unrolling_cfg - L = cfg.segments - if not 0.0 <= cfg.momentum < 1.0: - raise ValueError(f"momentum must be in [0, 1), got {cfg.momentum}.") - momentum_scale = 1.0 / (1.0 - cfg.momentum) ->>>>>>> 45df0de8 (feat: ADAM SOURCE and per-checkpoint optimizer saving in Trainer) if cfg.lr_list and cfg.step_size_list: return [ lr * k * momentum_scale for lr, k in zip(cfg.lr_list, cfg.step_size_list) diff --git a/bergson/approx_unrolling/pipeline.py b/bergson/approx_unrolling/pipeline.py index 99f1b2d5..971ea337 100644 --- a/bergson/approx_unrolling/pipeline.py +++ b/bergson/approx_unrolling/pipeline.py @@ -115,9 +115,8 @@ def approx_unrolling_pipeline( if not state_path.exists(): raise FileNotFoundError( f"use_adam_preconditioner requires {state_path}; write it " - "during training with TrainingConfig.save_optimizer_state, " - "then place it with bergson.utils.optimizer_placement." - "place_optimizer_states(save_dir, checkpoints)." + "during training with save_optimizer_state='all', then " + "export with bergson.utils.trainer_export.export_checkpoints." ) lr_times_steps_per_segment = compute_lr_times_steps_per_segment( diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index c2041a75..e6a134da 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -253,17 +253,12 @@ def worker( fsdp=run_cfg.fsdp, max_grad_norm=run_cfg.max_grad_norm, grad_accum_steps=run_cfg.grad_accum_steps, -<<<<<<< HEAD optimizer_cfg=( -======= - optimizer_hparams=( ->>>>>>> 5139d31a (refactor: clarify ADAM SOURCE naming; cover the applicator wiring) dict( betas=(run_cfg.adam_beta1, run_cfg.adam_beta2), eps=run_cfg.adam_eps, eps_root=run_cfg.eps_root, ) -<<<<<<< HEAD if getattr(run_cfg, "save_optimizer_state", "none") == "all" else None ), @@ -271,13 +266,6 @@ def worker( # Called on every rank: FSDP moments are DTensors whose gather is a # collective; rank 0 writes inside. if getattr(run_cfg, "save_optimizer_state", "none") != "none": -======= - if run_cfg.save_optimizer_state - else None - ), - ) - if run_cfg.save_optimizer_state: ->>>>>>> 5139d31a (refactor: clarify ADAM SOURCE naming; cover the applicator wiring) save_second_moments_as_optimizer_pt( model, fwd_state.opt_state, diff --git a/bergson/utils/optimizer_placement.py b/bergson/utils/optimizer_placement.py deleted file mode 100644 index 17f202b8..00000000 --- a/bergson/utils/optimizer_placement.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Place trainer-written optimizer states where attribution methods look for them. - -The trainer writes each snapshot's second moments as a *sibling* file, -``/step_.optimizer.pt``, next to that step's ``step_.ckpt`` -(see :meth:`bergson.magic.trainer.Trainer.train`). Both attribution consumers -instead expect ``optimizer.pt`` *inside* a directory: - -- SOURCE / approximate unrolling reads ``/optimizer.pt`` for every - checkpoint when ``ApproxUnrollingConfig.use_adam_preconditioner`` is set. -- TrackStar (and any ``build``) resolves ``PreprocessConfig.optimizer_state`` - via :func:`bergson.utils.load_from_optimizer.load_optimizer`, which accepts a - directory and loads ``optimizer.pt`` from inside it. - -Exporting a run to HF-format checkpoint dirs does not carry the sibling files -across, which is what previously forced a manual copy. :func:`place_optimizer_states` -does that placement, and satisfies both consumers with one call since they agree -on the ``/optimizer.pt`` layout. -""" - -import os -import re -import shutil -from pathlib import Path -from typing import Literal, Sequence - -OPTIMIZER_STATE_FILE = "optimizer.pt" -"""Filename both SOURCE and TrackStar look for inside a checkpoint directory.""" - -_STEP_OPTIMIZER_RE = re.compile(r"^step_(\d+)\.optimizer\.pt$") -_STEP_DIR_RE = re.compile(r"step_(\d+)") - -PlacementMode = Literal["copy", "symlink", "hardlink", "move"] - - -def sorted_optimizer_states(save_dir: str | Path) -> list[tuple[int, Path]]: - """Return ``(step, path)`` for each ``step_.optimizer.pt``, sorted by step. - - Mirrors :func:`bergson.data.sorted_checkpoints`, which does the same for the - ``step_.ckpt`` directories these files sit beside. - """ - save_dir = Path(save_dir) - if not save_dir.is_dir(): - raise NotADirectoryError(f"{save_dir} is not a directory") - - found = [] - for name in os.listdir(save_dir): - match = _STEP_OPTIMIZER_RE.match(name) - if match: - found.append((int(match.group(1)), save_dir / name)) - - return sorted(found, key=lambda pair: pair[0]) - - -def _step_of(path: Path) -> int | None: - """Step index named by a checkpoint directory, or None if it names none.""" - match = _STEP_DIR_RE.search(path.name) - return int(match.group(1)) if match else None - - -def _place(src: Path, dst: Path, mode: PlacementMode) -> None: - dst.parent.mkdir(parents=True, exist_ok=True) - if dst.exists() or dst.is_symlink(): - dst.unlink() - - match mode: - case "copy": - shutil.copy2(src, dst) - case "move": - shutil.move(str(src), str(dst)) - case "symlink": - # Relative, so an exported run stays valid if the tree is moved. - dst.symlink_to(os.path.relpath(src, dst.parent)) - case "hardlink": - os.link(src, dst) - case other: - raise ValueError(f"Unsupported placement mode: {other}") - - -def place_optimizer_states( - save_dir: str | Path, - checkpoints: Sequence[str | Path], - *, - mode: PlacementMode = "symlink", - overwrite: bool = False, -) -> dict[Path, Path]: - """Put an ``optimizer.pt`` inside each checkpoint dir, from ``save_dir``. - - ``save_dir`` is the trainer's output directory, holding the - ``step_.optimizer.pt`` files written when - ``TrainingConfig.save_optimizer_state`` is set. ``checkpoints`` are the - destination directories -- typically the exported HF checkpoints listed in - ``ApproxUnrollingConfig.checkpoints``. - - States are matched to destinations by step index when the destination - directory names one (``step_12`` / ``step_12.ckpt``); otherwise they are - matched positionally in step order, which requires the counts to be equal. - Mixing the two is an error, since a partial name match usually means the - caller passed a directory list that does not correspond to this run. - - ``mode`` defaults to ``"symlink"``: these files are large (second moments - are one scalar per weight) and duplicating them per checkpoint is wasteful. - Use ``"copy"`` when the destinations must be self-contained, e.g. before - uploading them somewhere. - - Returns ``{destination optimizer.pt: source file}``. Raises rather than - silently skipping, so a wrong directory list fails loudly instead of - producing a preconditioner from the wrong steps. - """ - states = sorted_optimizer_states(save_dir) - if not states: - raise FileNotFoundError( - f"No step_.optimizer.pt files in {save_dir}. Train with " - "TrainingConfig.save_optimizer_state=True to write them." - ) - - dsts = [Path(c) for c in checkpoints] - if not dsts: - raise ValueError("checkpoints is empty; nothing to place") - - missing = [d for d in dsts if not d.is_dir()] - if missing: - raise NotADirectoryError( - f"Checkpoint directories do not exist: {[str(m) for m in missing]}" - ) - - named = [_step_of(d) for d in dsts] - if all(step is not None for step in named): - steps = [step for step in named if step is not None] - by_step = dict(states) - unknown = [step for step in steps if step not in by_step] - if unknown: - raise FileNotFoundError( - f"No step_.optimizer.pt in {save_dir} for steps {unknown}; " - f"available steps: {sorted(by_step)}" - ) - pairs = [(by_step[step], dst) for step, dst in zip(steps, dsts)] - elif any(step is not None for step in named): - raise ValueError( - "Some checkpoint directories name a step and some do not, so they " - "cannot be matched reliably; pass directories that all name their " - "step, or none that do (for positional matching)." - ) - else: - if len(states) != len(dsts): - raise ValueError( - f"Cannot match positionally: {len(states)} optimizer states in " - f"{save_dir} but {len(dsts)} checkpoint directories. Name the " - "destination dirs after their steps to match by step instead." - ) - pairs = [(src, dst) for (_, src), dst in zip(states, dsts)] - - if not overwrite: - existing = [dst / OPTIMIZER_STATE_FILE for _, dst in pairs] - clashes = [p for p in existing if p.exists() or p.is_symlink()] - if clashes: - raise FileExistsError( - f"{[str(c) for c in clashes]} already exist; pass overwrite=True " - "to replace them." - ) - - placed: dict[Path, Path] = {} - for src, dst in pairs: - target = dst / OPTIMIZER_STATE_FILE - _place(src, target, mode) - placed[target] = src - - return placed - - -def place_final_optimizer_state( - run_path: str | Path, - destination: str | Path, - *, - mode: PlacementMode = "symlink", - overwrite: bool = False, -) -> Path: - """Put the run's final ``optimizer.pt`` inside ``destination``. - - ``TrainingConfig.save_optimizer_state`` writes the end-of-training state to - ``/optimizer.pt``. This places it inside an exported model - directory so ``PreprocessConfig.optimizer_state`` can point at that - directory -- the TrackStar gradient-normalization path. - - Returns the path written. - """ - src = Path(run_path) / OPTIMIZER_STATE_FILE - if not src.is_file(): - raise FileNotFoundError( - f"{src} not found. Train with TrainingConfig.save_optimizer_state=True " - "to write it." - ) - - dst_dir = Path(destination) - if not dst_dir.is_dir(): - raise NotADirectoryError(f"{dst_dir} is not a directory") - - target = dst_dir / OPTIMIZER_STATE_FILE - if target.resolve() == src.resolve(): - return target - if not overwrite and (target.exists() or target.is_symlink()): - raise FileExistsError(f"{target} already exists; pass overwrite=True.") - - _place(src, target, mode) - return target diff --git a/examples/bank_baselines/README.md b/examples/bank_baselines/README.md deleted file mode 100644 index 98765241..00000000 --- a/examples/bank_baselines/README.md +++ /dev/null @@ -1,41 +0,0 @@ -# Non-gradient attribution baselines - -Text similarity baselines evaluated by mean LDS over 50 test queries. - -- `bm25_baseline.py` — BM25 lexical overlap: pure surface-form term overlap, no model or embedding. -- `gradient_baseline.py` — gradient cosine similarity: cosine of the full per-example loss gradients on the bank's model (TracIn-style, unpreconditioned). -- `activation_baseline.py` — activation similarity: each doc is the mean-pooled input activation to every linear matrix of the model, L2-normalized per matrix and concatenated, then cosine similarity. -- `semantic_baseline.py` — semantic search with `jinaai/jina-embeddings-v3` (asymmetric `retrieval.query`/`retrieval.passage`). -- `qwen3_baseline.py` — semantic search with `Qwen/Qwen3-Embedding-8B`, a SOTA decoder embedder (`--model` to swap). - -Each produces a `[num_train_docs, num_queries]` score matrix. - -## Running - -Point a baseline at a re-train bank (written with `save_retrained_models=true`); it reads the model and dataset from the bank's `config.yaml`: - -```bash -python -m examples.bank_baselines.bm25_baseline --bank runs/retrain_bank_path -python -m examples.bank_baselines.gradient_baseline --bank runs/retrain_bank_path -python -m examples.bank_baselines.activation_baseline --bank runs/retrain_bank_path -python -m examples.bank_baselines.semantic_baseline --bank runs/retrain_bank_path -python -m examples.bank_baselines.qwen3_baseline --bank runs/retrain_bank_path -``` - -Omit `--bank` to build the default GPT-2/WikiText bank (`examples/magic/gpt2_wikitext_bank.yaml`) first. `--query_split` sets the query set (default `test[1:51]`), `--out` the output dir (default `runs/bank_baselines/`). - -## Results - -GPT-2 / WikiText bank (100 subsets, 1% leave-out, `eps_root 1e-8`): - -| Method | mean ρ | -| --- | --- | -| BM25 lexical overlap | 0.16 | -| Qwen3-Embedding-8B | 0.11 | -| activation similarity | 0.09 | -| Jina v3 semantic search | 0.06 | -| gradient cosine similarity | 0.05 | - -## Notes - -`jina-embeddings-v3`'s custom code predates transformers 5.x; `semantic_baseline.load_model` sets an `all_tied_weights_keys` default and resets the NaN LoRA `lora_dropout_mask` buffers to ones so it loads and runs. NVIDIA's NV-Embed-v2 is a comparable SOTA embedder but its custom code is incompatible with transformers 5.x, so `qwen3_baseline.py` uses Qwen3-Embedding instead. diff --git a/examples/bank_baselines/__init__.py b/examples/bank_baselines/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/examples/bank_baselines/activation_baseline.py b/examples/bank_baselines/activation_baseline.py deleted file mode 100644 index ac99fb0e..00000000 --- a/examples/bank_baselines/activation_baseline.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Baseline: activation similarity over the attribution target modules. - -Represents each document by the mean-pooled input activation to every target -linear module of the bank's fully-trained model (the same modules the gradient -methods attribute through: all Conv1D/Linear except ``lm_head``). Each -per-module pooled vector is L2-normalized, then concatenated, so cosine -similarity between two docs equals the average per-module cosine similarity. - -A training doc that is activation-similar to a query is predicted to be -influential (removing it should raise query loss), so the loss-diff-convention -score is ``-cosine`` -- larger => larger predicted loss reduction from keeping -the doc. - -Run with (builds the default bank if --bank is omitted): - python -m examples.bank_baselines.activation_baseline --bank runs/retrain_bank_path -""" - -import argparse -from pathlib import Path - -import numpy as np -import torch -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.pytorch_utils import Conv1D - -from . import common - - -def target_module_names(model: torch.nn.Module) -> list[str]: - """Linear/Conv1D modules the gradient methods attribute through (not lm_head).""" - names = [] - for name, mod in model.named_modules(): - if name.endswith("lm_head"): - continue - if isinstance(mod, (Conv1D, torch.nn.Linear)): - names.append(name) - return names - - -@torch.no_grad() -def embed( - texts: list[str], - model: torch.nn.Module, - tokenizer, - module_names: list[str], - device: str, - max_len: int, - batch_size: int, -) -> np.ndarray: - """Per-module-normalized, concatenated mean-pooled activation embeddings.""" - modules = dict(model.named_modules()) - captured: dict[str, torch.Tensor] = {} - handles = [] - for mn in module_names: - - def hook(_m, inp, _out, _mn=mn): - captured[_mn] = inp[0].detach() - - handles.append(modules[mn].register_forward_hook(hook)) - - out_rows = [] - try: - for start in range(0, len(texts), batch_size): - batch = texts[start : start + batch_size] - enc = tokenizer( - batch, - return_tensors="pt", - padding=True, - truncation=True, - max_length=max_len, - ).to(device) - mask = enc["attention_mask"].unsqueeze(-1).float() # [B, T, 1] - captured.clear() - model(**enc) - - per_module = [] - denom = mask.sum(dim=1).clamp_min(1.0) # [B, 1] - for mn in module_names: - act = captured[mn].float() # [B, T, d] - pooled = (act * mask).sum(dim=1) / denom # [B, d] - pooled = torch.nn.functional.normalize(pooled, dim=-1) - per_module.append(pooled) - emb = torch.cat(per_module, dim=-1) # [B, sum_d] - out_rows.append(emb.cpu().numpy()) - finally: - for h in handles: - h.remove() - return np.concatenate(out_rows, axis=0) - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") - ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) - ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) - ap.add_argument("--device", default="cuda:0") - ap.add_argument("--max_len", type=int, default=1024) - ap.add_argument("--batch_size", type=int, default=32) - args = ap.parse_args() - - bank = common.ensure_bank(args.bank) - spec = common.read_bank_spec(bank) - out_dir = Path(args.out) - - tokenizer = AutoTokenizer.from_pretrained(spec.model) - tokenizer.pad_token = tokenizer.eos_token - model = AutoModelForCausalLM.from_pretrained( - spec.base_model, dtype=torch.float32, attn_implementation="eager" - ).to(args.device) - model.eval() - - module_names = target_module_names(model) - print(f"Embedding over {len(module_names)} modules, e.g. {module_names[:3]} ...") - - train_texts, query_texts = common.load_texts(spec, args.query_split) - print(f"Embedding {len(train_texts)} train docs ...") - train_emb = embed( - train_texts, - model, - tokenizer, - module_names, - args.device, - args.max_len, - args.batch_size, - ) - print(f"Embedding {len(query_texts)} query docs ...") - query_emb = embed( - query_texts, - model, - tokenizer, - module_names, - args.device, - args.max_len, - args.batch_size, - ) - - # Cosine similarity (rows already unit-norm per module; renormalize the - # concatenation so cosine == mean per-module cosine). - tn = train_emb / np.linalg.norm(train_emb, axis=1, keepdims=True) - qn = query_emb / np.linalg.norm(query_emb, axis=1, keepdims=True) - cosine = tn @ qn.T # [num_train, num_query] - - scores = -cosine # loss-diff convention (similar => influential) - score_path = common.save_scores(scores, out_dir, "activation_scores") - - rhos = common.evaluate_lds( - bank, score_path, out_dir / "activation_validate", spec, args.query_split - ) - common.report("activation similarity", rhos) - - -if __name__ == "__main__": - main() diff --git a/examples/bank_baselines/bm25_baseline.py b/examples/bank_baselines/bm25_baseline.py deleted file mode 100644 index cf7f8928..00000000 --- a/examples/bank_baselines/bm25_baseline.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Baseline: BM25 lexical overlap. - -Scores each training doc against each query by Okapi BM25 -- pure surface-form -term overlap, no model and no learned embedding. A lexical proxy for influence: -docs sharing query terms are predicted influential, so the loss-diff-convention -score is ``-bm25``. Often a stronger influence proxy than deep-semantic -embedders for small LMs, and essentially free. - -Run with (builds the default bank if --bank is omitted): - python -m examples.bank_baselines.bm25_baseline --bank runs/retrain_bank_path -""" - -import argparse -from pathlib import Path - -import numpy as np -from sklearn.feature_extraction.text import CountVectorizer - -from . import common - - -def bm25_scores( - train_texts: list[str], query_texts: list[str], k1: float = 1.5, b: float = 0.75 -) -> np.ndarray: - """Okapi BM25 score of every train doc against every query -> [docs, queries].""" - vectorizer = CountVectorizer() - tf = vectorizer.fit_transform(train_texts).astype(np.float64).tocsr() # [D, V] - n_docs, _ = tf.shape - - doc_len = np.asarray(tf.sum(axis=1)).ravel() - avgdl = doc_len.mean() - df = np.asarray((tf > 0).sum(axis=0)).ravel() - idf = np.log((n_docs - df + 0.5) / (df + 0.5) + 1.0) - - # Per-nonzero BM25 term weight: idf * tf*(k1+1) / (tf + k1*(1-b+b*dl/avgdl)). - rows, cols = tf.nonzero() - denom_bias = k1 * (1 - b + b * doc_len[rows] / avgdl) - weighted = tf.copy() - weighted.data = idf[cols] * tf.data * (k1 + 1.0) / (tf.data + denom_bias) - - # Query terms as a binary term set (Okapi BM25 ignores query term frequency). - q_bin = (vectorizer.transform(query_texts) > 0).astype(np.float64) # [Q, V] - return np.asarray((weighted @ q_bin.T).todense()) # [D, Q] - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") - ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) - ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) - args = ap.parse_args() - - bank = common.ensure_bank(args.bank) - spec = common.read_bank_spec(bank) - out_dir = Path(args.out) - - train_texts, query_texts = common.load_texts(spec, args.query_split) - print(f"BM25 over {len(train_texts)} train docs, {len(query_texts)} queries ...") - scores = -bm25_scores(train_texts, query_texts) # loss-diff convention - score_path = common.save_scores(scores, out_dir, "bm25_scores") - - rhos = common.evaluate_lds( - bank, score_path, out_dir / "bm25_validate", spec, args.query_split - ) - common.report("BM25 lexical overlap", rhos) - - -if __name__ == "__main__": - main() diff --git a/examples/bank_baselines/common.py b/examples/bank_baselines/common.py deleted file mode 100644 index 2da7038b..00000000 --- a/examples/bank_baselines/common.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Shared helpers for non-gradient attribution baselines on a re-train bank. - -A "re-train bank" is any directory written with ``save_retrained_models=true`` -(e.g. by ``examples/magic/gpt2_wikitext_bank.yaml``): it holds ``subsets.json``, -``retrained/base`` and ``retrained/subset_*`` checkpoints, and the ``config.yaml`` -that built it. Each baseline reads the model + dataset straight from the bank, -produces a ``[num_train_docs, num_queries]`` score matrix, and evaluates it by -per-query LDS -- so ``--bank `` is all anyone needs to run one. - -The LDS itself reuses ``bergson``'s ``evaluate_retrained``: it computes each -subset's query-loss diff from the banked models once, caches it under the bank, -and every later baseline on the same bank/query reuses that cache. -""" - -import subprocess -from dataclasses import dataclass -from pathlib import Path - -import numpy as np -import yaml -from datasets import load_dataset - -from bergson.config.config import DataConfig -from bergson.magic.config import MagicConfig -from bergson.validate import evaluate_retrained - -REPO = Path(__file__).resolve().parents[2] - -# The default bank built on demand when ``--bank`` is omitted. -DEFAULT_BANK = REPO / "runs" / "gpt2_wikitext_bank" -DEFAULT_BANK_CONFIG = REPO / "examples" / "magic" / "gpt2_wikitext_bank.yaml" -# LDS query set (50 held-out docs); overridable per run. -DEFAULT_QUERY_SPLIT = "test[1:51]" - - -@dataclass -class BankSpec: - """The model + data a bank was built on, read from its ``config.yaml``.""" - - model: str - base_model: str # retrained/base if present, else the base model id - dataset: str - train_split: str - prompt_column: str - chunk_length: int - batch_size: int - - -def ensure_bank(bank: str | None) -> Path: - """Return a ready re-train bank, building the default one if none is given. - - A bank is ready when it has ``subsets.json`` and ``retrained/base``. When - ``bank`` is ``None`` and the default bank is absent, it is built by running - its config through ``bergson`` (a leave-k-out re-training job). - """ - if bank is not None: - path = Path(bank) - if not (path / "subsets.json").exists(): - raise FileNotFoundError( - f"{path} is not a re-train bank (no subsets.json); pass a dir " - "written with save_retrained_models=true" - ) - return path - - if not (DEFAULT_BANK / "subsets.json").exists(): - cmd = ["bergson", str(DEFAULT_BANK_CONFIG)] - print( - f"No bank passed and {DEFAULT_BANK} missing; building it:\n" - f" {' '.join(cmd)}" - ) - subprocess.run(cmd, cwd=REPO, check=True) - return DEFAULT_BANK - - -def read_bank_spec(bank: Path) -> BankSpec: - """Read the model + dataset a bank was built on from its ``config.yaml``.""" - with open(bank / "config.yaml") as f: - config = yaml.safe_load(f) - # The bank-building step (magic/train) is the one carrying model + data. - step = next( - v - for s in config["steps"] - for v in s.values() - if isinstance(v, dict) and "model" in v and "data" in v - ) - data = step["data"] - base = bank / "retrained" / "base" - return BankSpec( - model=step["model"], - base_model=str(base) if base.exists() else step["model"], - dataset=data["dataset"], - train_split=data.get("split", "train"), - prompt_column=data.get("prompt_column", "text"), - chunk_length=data.get("chunk_length", 0), - batch_size=step.get("batch_size", 64), - ) - - -def load_texts(spec: BankSpec, query_split: str) -> tuple[list[str], list[str]]: - """Return (train_texts, query_texts) in original doc order. - - Row i of ``train_texts`` is training doc id i, matching the ids in the - bank's ``subsets.json`` and the rows of every score matrix. - """ - train = load_dataset(spec.dataset, split=spec.train_split) - query = load_dataset(spec.dataset, split=query_split) - return list(train[spec.prompt_column]), list(query[spec.prompt_column]) - - -def evaluate_lds( - bank: Path, - score_path: Path, - out_dir: Path, - spec: BankSpec, - query_split: str, -) -> np.ndarray: - """Per-query LDS Spearman of a score matrix against the bank. - - Runs ``evaluate_retrained`` (no re-training; reuses the bank's cached - per-subset query losses) and returns the per-query Spearman correlations it - writes to ``summary.csv``. - """ - run_cfg = MagicConfig( - run_path=str(out_dir), - model=spec.model, - precision="fp32", - batch_size=spec.batch_size, - query=DataConfig( - dataset=spec.dataset, - split=query_split, - prompt_column=spec.prompt_column, - chunk_length=0, - ), - ) - evaluate_retrained(run_cfg, str(bank), score_path=str(score_path)) - - summary = np.genfromtxt(out_dir / "summary.csv", delimiter=",", names=True) - return np.atleast_1d(summary["spearman_corr"]).astype(float) - - -def save_scores(scores: np.ndarray, out_dir: Path, name: str) -> Path: - """Save a ``[num_train_docs, num_queries]`` score matrix as ``.npy``.""" - out_dir.mkdir(parents=True, exist_ok=True) - path = out_dir / f"{name}.npy" - np.save(path, scores.astype(np.float32)) - print(f"Saved scores {scores.shape} -> {path}") - return path - - -def report(name: str, rhos: np.ndarray) -> None: - """Print a per-query LDS summary consistent with the repo's convention.""" - print(f"\n=== {name}: per-query LDS Spearman (n={len(rhos)} queries) ===") - print(f" mean {np.mean(rhos):+.4f}") - print(f" median {np.median(rhos):+.4f}") - print(f" min {np.min(rhos):+.4f} max {np.max(rhos):+.4f}") - print(f" frac>0 {np.mean(rhos > 0):.2f}") diff --git a/examples/bank_baselines/gradient_baseline.py b/examples/bank_baselines/gradient_baseline.py deleted file mode 100644 index 23416113..00000000 --- a/examples/bank_baselines/gradient_baseline.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Baseline: gradient cosine similarity (TracIn-style, full parameters). - -Scores a train doc against a query by the cosine similarity of their full -per-example loss gradients on the bank's model -- the raw, unpreconditioned -influence signal that TrackStar/EK-FAC refine. A train doc whose gradient -aligns with the query's is predicted influential, so the loss-diff-convention -score is ``-cosine``. - -Full gradients are never stored: the 50 query gradients stay resident on the -GPU (GPT-2 is 124M params, ~25 GB in fp32) and each train-doc gradient is -computed, dotted against all queries, and discarded. - -Run with (builds the default bank if --bank is omitted): - python -m examples.bank_baselines.gradient_baseline --bank runs/retrain_bank_path -""" - -import argparse -from pathlib import Path - -import numpy as np -import torch -from transformers import AutoModelForCausalLM, AutoTokenizer - -from . import common - - -def per_example_grad( - text: str, model, tokenizer, device: str, max_len: int -) -> torch.Tensor: - """Flattened full-parameter gradient of the mean-CE loss on one document.""" - enc = tokenizer(text, return_tensors="pt", truncation=True, max_length=max_len).to( - device - ) - model.zero_grad(set_to_none=True) - model(input_ids=enc["input_ids"], labels=enc["input_ids"]).loss.backward() - return torch.cat( - [ - (p.grad if p.grad is not None else torch.zeros_like(p)).flatten() - for p in model.parameters() - ] - ) - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") - ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) - ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) - ap.add_argument("--device", default="cuda:0") - ap.add_argument("--max_len", type=int, default=1024) - args = ap.parse_args() - - bank = common.ensure_bank(args.bank) - spec = common.read_bank_spec(bank) - out_dir = Path(args.out) - - tokenizer = AutoTokenizer.from_pretrained(spec.model) - model = AutoModelForCausalLM.from_pretrained( - spec.base_model, dtype=torch.float32, attn_implementation="eager" - ).to(args.device) - model.eval() - - train_texts, query_texts = common.load_texts(spec, args.query_split) - - # Keep the query gradients resident; stream train gradients against them. - # Pre-allocate and fill row by row so we never hold a second copy (each full - # gradient is ~0.5 GB, the [Q, P] block ~25 GB in fp32). - print(f"Computing {len(query_texts)} query gradients ...") - n_params = sum(p.numel() for p in model.parameters()) - query_grads = torch.empty(len(query_texts), n_params, device=args.device) - for i, t in enumerate(query_texts): - query_grads[i] = per_example_grad(t, model, tokenizer, args.device, args.max_len) - query_norms = query_grads.norm(dim=1).clamp_min(1e-12) - - print(f"Scoring {len(train_texts)} train gradients against queries ...") - scores = np.empty((len(train_texts), len(query_texts)), dtype=np.float32) - for i, text in enumerate(train_texts): - g = per_example_grad(text, model, tokenizer, args.device, args.max_len) - cos = (query_grads @ g) / (query_norms * g.norm().clamp_min(1e-12)) - scores[i] = cos.detach().cpu().numpy() - - scores = -scores # loss-diff convention (aligned gradient => influential) - score_path = common.save_scores(scores, out_dir, "gradient_scores") - - rhos = common.evaluate_lds( - bank, score_path, out_dir / "gradient_validate", spec, args.query_split - ) - common.report("gradient cosine similarity", rhos) - - -if __name__ == "__main__": - main() diff --git a/examples/bank_baselines/qwen3_baseline.py b/examples/bank_baselines/qwen3_baseline.py deleted file mode 100644 index 42576ace..00000000 --- a/examples/bank_baselines/qwen3_baseline.py +++ /dev/null @@ -1,76 +0,0 @@ -"""Baseline: semantic search with Qwen3-Embedding-8B. - -A SOTA decoder-based embedding model (causal-attention Qwen3 backbone, -last-token pooling), top of the open MTEB leaderboard -- a much stronger, -attention-based embedder than the jina-v3 encoder. Training docs are embedded -as passages and queries with the retrieval prompt; a train doc scores against a -query by cosine similarity, so the loss-diff-convention score is ``-cosine``. - -(NVIDIA's NV-Embed-v2 is a comparable model but its custom modeling code is -incompatible with transformers 5.x; Qwen3-Embedding uses the native Qwen3 -architecture and loads cleanly. Pass --model Qwen/Qwen3-Embedding-4B for the -smaller, faster variant.) - -Run with (builds the default bank if --bank is omitted): - python -m examples.bank_baselines.qwen3_baseline --bank runs/retrain_bank_path -""" - -import argparse -from pathlib import Path - -from sentence_transformers import SentenceTransformer - -from . import common - -MODEL = "Qwen/Qwen3-Embedding-8B" - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") - ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) - ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) - ap.add_argument("--model", default=MODEL) - ap.add_argument("--device", default="cuda:0") - ap.add_argument("--max_length", type=int, default=512) - ap.add_argument("--batch_size", type=int, default=8) - args = ap.parse_args() - - bank = common.ensure_bank(args.bank) - spec = common.read_bank_spec(bank) - out_dir = Path(args.out) - - model = SentenceTransformer( - args.model, model_kwargs={"torch_dtype": "bfloat16"}, device=args.device - ) - model.max_seq_length = args.max_length - - train_texts, query_texts = common.load_texts(spec, args.query_split) - print(f"Embedding {len(train_texts)} train docs (passages) ...") - train_emb = model.encode( - train_texts, - batch_size=args.batch_size, - normalize_embeddings=True, - show_progress_bar=True, - ) - print(f"Embedding {len(query_texts)} query docs (query prompt) ...") - query_emb = model.encode( - query_texts, - prompt_name="query", - batch_size=args.batch_size, - normalize_embeddings=True, - show_progress_bar=True, - ) - - cosine = train_emb @ query_emb.T # rows unit-norm => dot == cosine - scores = -cosine # loss-diff convention (similar => influential) - score_path = common.save_scores(scores, out_dir, "qwen3_scores") - - rhos = common.evaluate_lds( - bank, score_path, out_dir / "qwen3_validate", spec, args.query_split - ) - common.report(f"{args.model} semantic similarity", rhos) - - -if __name__ == "__main__": - main() diff --git a/examples/bank_baselines/semantic_baseline.py b/examples/bank_baselines/semantic_baseline.py deleted file mode 100644 index 2103251d..00000000 --- a/examples/bank_baselines/semantic_baseline.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Baseline: semantic-search similarity with a Jina AI embedding model. - -Embeds each training document as a retrieval passage and each query document -as a retrieval query with ``jinaai/jina-embeddings-v3`` (a strong 570M-param, -1024-dim, 8192-context text embedder), then scores a train doc against a query -by their cosine similarity -- ordinary dense semantic search. This ignores the -attributed model entirely; it is a pure content-similarity baseline. - -A training doc semantically similar to a query is predicted to be influential -(removing it should raise query loss), so the loss-diff-convention score is -``-cosine``. - -jina-embeddings-v3 ships custom modeling code that predates transformers 5.x, -so ``load_model`` patches two load-time incompatibilities (see there) to run it -for inference. - -Run with (builds the default bank if --bank is omitted): - python -m examples.bank_baselines.semantic_baseline --bank runs/retrain_bank_path -""" - -import argparse -from pathlib import Path - -import numpy as np -import torch -from transformers.modeling_utils import PreTrainedModel - -from . import common - -MODEL = "jinaai/jina-embeddings-v3" - - -def load_model(device: str): - # jina-embeddings-v3's custom code predates transformers 5.x; two fixes: - # (1) from_pretrained reads all_tied_weights_keys, which the custom class - # never defines -- give a benign empty default so load doesn't crash. - if not isinstance( - getattr(PreTrainedModel, "all_tied_weights_keys", None), property - ): - PreTrainedModel.all_tied_weights_keys = {} - from transformers import AutoModel - - model = AutoModel.from_pretrained( - MODEL, trust_remote_code=True, dtype=torch.float32 - ) - - # (2) its LoRA task adapters leave the per-forward lora_dropout_mask buffers - # uninitialized (NaN), so every embedding comes out NaN. In eval there - # is no dropout, so reset them to ones. - for name, buf in model.named_buffers(): - if "lora_dropout_mask" in name: - buf.data = torch.ones_like(buf) - return model.to(device).eval() - - -@torch.no_grad() -def encode(model, texts: list[str], task: str, batch_size: int) -> np.ndarray: - """L2-normalized embeddings via jina's task-specific encode API.""" - out = [] - for start in range(0, len(texts), batch_size): - emb = model.encode( - texts[start : start + batch_size], - task=task, - convert_to_numpy=True, - normalize_embeddings=True, - ) - out.append(np.asarray(emb, dtype=np.float32)) - return np.concatenate(out, axis=0) - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") - ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) - ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) - ap.add_argument("--device", default="cuda:0") - ap.add_argument("--batch_size", type=int, default=16) - args = ap.parse_args() - - bank = common.ensure_bank(args.bank) - spec = common.read_bank_spec(bank) - out_dir = Path(args.out) - model = load_model(args.device) - - train_texts, query_texts = common.load_texts(spec, args.query_split) - print(f"Embedding {len(train_texts)} train docs (retrieval.passage) ...") - train_emb = encode(model, train_texts, "retrieval.passage", args.batch_size) - print(f"Embedding {len(query_texts)} query docs (retrieval.query) ...") - query_emb = encode(model, query_texts, "retrieval.query", args.batch_size) - - cosine = train_emb @ query_emb.T # rows unit-norm => dot == cosine - scores = -cosine # loss-diff convention (similar => influential) - score_path = common.save_scores(scores, out_dir, "semantic_scores") - - rhos = common.evaluate_lds( - bank, score_path, out_dir / "semantic_validate", spec, args.query_split - ) - common.report("Jina v3 semantic similarity", rhos) - - -if __name__ == "__main__": - main() diff --git a/examples/magic/gpt2_wikitext_bank.yaml b/examples/magic/gpt2_wikitext_bank.yaml deleted file mode 100644 index f47758b4..00000000 --- a/examples/magic/gpt2_wikitext_bank.yaml +++ /dev/null @@ -1,43 +0,0 @@ -# Build a GPT-2 / WikiText leave-k-out re-train bank (single node): MAGIC -# attribution plus 400 retrained leave-out models saved for later evaluation. - -# Run with `bergson examples/magic/gpt2_wikitext_bank.yaml` - -steps: - - magic: - run_path: runs/gpt2_wikitext_bank - model: gpt2 - overwrite: true - cleanup_ckpts: false - - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 0 - - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[0:1]" - chunk_length: 0 - - distributed: - nproc_per_node: 8 - nnode: 1 - - eps_root: 1e-8 - - batch_size: 64 - num_epochs: 4 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: random - wandb_project: magic - num_subsets: 100 - subset_fraction: 0.01 - save_optimizer_state: True - save_retrained_models: True diff --git a/examples/magic/gpt2_wikitext_dropout.yaml b/examples/magic/gpt2_wikitext_dropout.yaml deleted file mode 100644 index 6c3c04a0..00000000 --- a/examples/magic/gpt2_wikitext_dropout.yaml +++ /dev/null @@ -1,40 +0,0 @@ -# LDS: 0.1862 - -steps: - - magic: - run_path: runs/gpt2_wikitext_dropout - model: gpt2 - overwrite: true - cleanup_ckpts: false - train_mode: true - weight_decay: 0.01 - adam_eps: 1.0e-8 - eps_root: 1.0e-6 - loss_reduction: mean - - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 0 - - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[0:1]" - chunk_length: 0 - - distributed: - nproc_per_node: 8 - nnode: 1 - - batch_size: 64 - num_epochs: 4 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - wandb_project: magic - save_optimizer_state: True - save_retrained_models: True diff --git a/examples/magic/gpt2_wikitext_dropout_s15.yaml b/examples/magic/gpt2_wikitext_dropout_s15.yaml deleted file mode 100644 index a9e7cd8c..00000000 --- a/examples/magic/gpt2_wikitext_dropout_s15.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# LDS: -0.2286 - -steps: - - magic: - run_path: runs/gpt2_wikitext_dropout_s15 - model: gpt2 - overwrite: true - cleanup_ckpts: false - train_mode: true - num_subsets: 15 - subset_fraction: 0.01 - weight_decay: 0.01 - adam_eps: 1.0e-8 - eps_root: 1.0e-6 - loss_reduction: mean - - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 0 - - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[0:1]" - chunk_length: 0 - - distributed: - nproc_per_node: 8 - nnode: 1 - - batch_size: 64 - num_epochs: 4 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - wandb_project: magic - save_optimizer_state: True - save_retrained_models: True diff --git a/examples/magic/gpt2_wikitext_lotus.yaml b/examples/magic/gpt2_wikitext_lotus.yaml deleted file mode 100644 index e1df5f3b..00000000 --- a/examples/magic/gpt2_wikitext_lotus.yaml +++ /dev/null @@ -1,44 +0,0 @@ -# Single-step MAGIC attribution example (GPT-2 / WikiText, single node). -# Isolated run_path ("lotus") so this doesn't collide with any other -# job sharing examples/magic/gpt2_wikitext.yaml's run_path. - -# Run with `bergson examples/magic/gpt2_wikitext_lotus.yaml` - -steps: - - magic: - run_path: runs/lotus - model: gpt2 - overwrite: true - cleanup_ckpts: false - - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 0 - - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[0:1]" - chunk_length: 0 - - distributed: - nproc_per_node: 8 - nnode: 1 - - eps_root: 1e-6 - - batch_size: 64 - num_epochs: 4 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: random - wandb_project: magic - num_subsets: 400 - subset_fraction: 0.01 - save_optimizer_state: True - save_retrained_models: True diff --git a/examples/magic/gpt2_wikitext_lotus_ekfac_q50.yaml b/examples/magic/gpt2_wikitext_lotus_ekfac_q50.yaml deleted file mode 100644 index e84af898..00000000 --- a/examples/magic/gpt2_wikitext_lotus_ekfac_q50.yaml +++ /dev/null @@ -1,79 +0,0 @@ -# EK-FAC counterpart to gpt2_wikitext_lotus_mq.yaml: computes EK-FAC -# attribution scores for the SAME base model trained by -# examples/magic/gpt2_wikitext_lotus.yaml (runs/lotus/retrained/base, -# bs=64/eps_root=1e-6), against the same 50 test queries (test[1:51]), -# then validates those scores against the existing 400-model leave-k-out -# bank in runs/lotus/retrained -- no retraining needed, that bank is -# already complete. -# -# IMPORTANT: run_path values below are meant to be unique on this shared -# filesystem. If another run is already using them, rename before running -# (e.g. add your hostname) -- do NOT run this against a run_path someone -# else might be using. runs/lotus itself is read-only here and already -# finished, so it's safe to read from concurrently. - -# Run with `bergson examples/magic/gpt2_wikitext_lotus_ekfac_q50.yaml` -# from /mnt/ssd-1/lucia/bergson-damping (paths below are relative to that). - -steps: - - ekfac: - index_cfg: - run_path: runs/lotus_scores_ekfac50q - model: runs/lotus/retrained/base - overwrite: true - # GPT-2's max context length is 1024. - token_batch_size: 1024 - # Cap batch size in docs so the bin-packer doesn't produce - # large batches of very few tokens (gradients aren't compressed - # for EK-FAC). - max_batch_size: 64 - # GPT-2 uses an nn.Linear for its LM head. - filter_modules: "lm_head" - distributed: - nproc_per_node: 8 # <-- adjust to this node's GPU count - nnode: 1 - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 0 - - hessian_cfg: - method: kfac - ev_correction: true - - score_cfg: - batch_size: 1024 - - preprocess_cfg: - unit_normalize: false - - hessian_pipeline_cfg: - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[1:51]" - chunk_length: 0 - # One score column per query (50 total), not aggregated -- - # required for the multi-query LDS validation below. - query_aggregation: none - inversion_cfg: - damping_factor: 0.1 - - - validate: - run_path: runs/lotus_scores_ekfac50q_validate - model: gpt2 - overwrite: true - - scores: runs/lotus_scores_ekfac50q/scores - retrained_dir: runs/lotus - - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 0 - - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[1:51]" - chunk_length: 0 - - batch_size: 64 diff --git a/examples/magic/gpt2_wikitext_lotus_final_q01_50.yaml b/examples/magic/gpt2_wikitext_lotus_final_q01_50.yaml deleted file mode 100644 index 26767b82..00000000 --- a/examples/magic/gpt2_wikitext_lotus_final_q01_50.yaml +++ /dev/null @@ -1,31 +0,0 @@ -# Final true multi-query LDS check on all 50 per-query MAGIC backward -# passes, recomputed post-fix by scratchpad/run_50_bwd_fixed.sh. No -# retraining, no backward pass -- reuses the existing 400-model leave-k-out -# bank at runs/lotus/retrained, scored against a stacked [4608, 50] -# multi-query score array (runs/lotus_bwd/final_q01_50_scores.npy). -# -# Run with `bergson examples/magic/gpt2_wikitext_lotus_final_q01_50.yaml` - -steps: - - validate: - run_path: runs/lotus_final_q01_50 - model: gpt2 - overwrite: true - - scores: runs/lotus_bwd/final_q01_50_scores.npy - retrained_dir: runs/lotus - - distributed: - nproc_per_node: 1 - - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 0 - - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[1:51]" - chunk_length: 0 - - batch_size: 8 diff --git a/examples/magic/gpt2_wikitext_lotus_mq.yaml b/examples/magic/gpt2_wikitext_lotus_mq.yaml deleted file mode 100644 index fe9f2bcb..00000000 --- a/examples/magic/gpt2_wikitext_lotus_mq.yaml +++ /dev/null @@ -1,70 +0,0 @@ -# Multi-query LDS follow-up to the "lotus" run (examples/magic/gpt2_wikitext_lotus.yaml). -# -# Step 1 retrains the same bs=64/eps_root=1e-6 base model (deterministic, -# same seed) purely to get a fresh MAGIC backward pass against 50 NEW test -# queries (test[1:51], i.e. the 50 WikiText test docs after the one used by -# the original lotus run) -- skip_validation stops it before any retraining. -# -# Step 2 reuses the 400 leave-k-out checkpoints already saved under -# runs/lotus/retrained (no retraining) and correlates them against the new -# per-query scores from step 1, reporting a mean Spearman LDS across the -# 50 queries. - -# Run with `bergson examples/magic/gpt2_wikitext_lotus_mq.yaml` - -steps: - - magic: - run_path: runs/lotus_mq_scores - model: gpt2 - overwrite: true - cleanup_ckpts: false - - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 0 - - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[1:51]" - chunk_length: 0 - - distributed: - nproc_per_node: 8 - nnode: 1 - - eps_root: 1e-6 - - batch_size: 64 - num_epochs: 4 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - skip_validation: true - save_retrained_models: false - save_optimizer_state: false - wandb_project: magic - - - validate: - run_path: runs/lotus_mq_eval - model: gpt2 - overwrite: true - - scores: runs/lotus_mq_scores/scores.pt - retrained_dir: runs/lotus - - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 0 - - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[1:51]" - chunk_length: 0 - - batch_size: 64 diff --git a/examples/magic/gpt2_wikitext_train_epsroot0_bs64.yaml b/examples/magic/gpt2_wikitext_train_epsroot0_bs64.yaml deleted file mode 100644 index 9aa8e6c0..00000000 --- a/examples/magic/gpt2_wikitext_train_epsroot0_bs64.yaml +++ /dev/null @@ -1,54 +0,0 @@ -steps: -- train: - run_path: runs/gpt2_wikitext_epsroot0_bs64 - overwrite: true - model: gpt2 - precision: fp32 - revision: null - distributed: - nnode: 1 - nproc_per_node: 8 - node_rank: null - fsdp: false - peft_init_kwargs: '' - model_kwargs: '' - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: train - subset: null - prompt_column: text - completion_column: '' - conversation_column: '' - reward_column: '' - skip_nan_rewards: false - truncation: false - format_template: '' - data_kwargs: '' - chunk_length: 0 - tokenizer: '' - drop_columns: true - max_tokens: null - use_tf32_matmuls: false - debug: false - lr_schedule: - lr: 0.0008 - lr_scheduler_type: polynomial - lr_start: 1.0e-06 - lr_end: 8.0e-05 - warmup_steps: 0.25 - num_cycles: 0.5 - power: 1.0 - batch_size: 64 - num_epochs: 4 - seed: 42 - adam_beta1: 0.95 - adam_beta2: 0.975 - eps_root: 0.0 - optimizer: adamw - save_optimizer_state: true - save_retrained_models: true - weight_decay: 0.01 - max_grad_norm: null - grad_checkpointing: false - resume: false - wandb_project: magic diff --git a/examples/magic/gpt2_wikitext_train_stdadam_bs64.yaml b/examples/magic/gpt2_wikitext_train_stdadam_bs64.yaml deleted file mode 100644 index 45abebb9..00000000 --- a/examples/magic/gpt2_wikitext_train_stdadam_bs64.yaml +++ /dev/null @@ -1,54 +0,0 @@ -steps: -- train: - run_path: runs/gpt2_wikitext_stdadam_bs64 - overwrite: true - model: gpt2 - precision: fp32 - revision: null - distributed: - nnode: 1 - nproc_per_node: 8 - node_rank: null - fsdp: false - peft_init_kwargs: '' - model_kwargs: '' - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: train - subset: null - prompt_column: text - completion_column: '' - conversation_column: '' - reward_column: '' - skip_nan_rewards: false - truncation: false - format_template: '' - data_kwargs: '' - chunk_length: 0 - tokenizer: '' - drop_columns: true - max_tokens: null - use_tf32_matmuls: false - debug: false - lr_schedule: - lr: 0.0008 - lr_scheduler_type: polynomial - lr_start: 1.0e-06 - lr_end: 8.0e-05 - warmup_steps: 0.25 - num_cycles: 0.5 - power: 1.0 - batch_size: 64 - num_epochs: 4 - seed: 42 - adam_beta1: 0.9 - adam_beta2: 0.999 - eps_root: 0.0 - optimizer: adamw - save_optimizer_state: true - save_retrained_models: true - weight_decay: 0.01 - max_grad_norm: null - grad_checkpointing: false - resume: false - wandb_project: magic diff --git a/examples/method_comparison/gpt2_wikitext_lotus_source_adam_q50.yaml b/examples/method_comparison/gpt2_wikitext_lotus_source_adam_q50.yaml deleted file mode 100644 index 467c03a1..00000000 --- a/examples/method_comparison/gpt2_wikitext_lotus_source_adam_q50.yaml +++ /dev/null @@ -1,91 +0,0 @@ -# Adam/AdamW SOURCE variant (Bae et al., 2024, Appendix C) for the lotus -# GPT-2 base model against the 50 test queries, validated against the -# complete 400-model lotus bank. The unrolling eigenfunctions are evaluated -# on a diagonal approximation of P^1/2 H P^1/2, where P is the per-segment -# diagonal Adam preconditioner built from each checkpoint dir's -# optimizer.pt (written per snapshot during training by -# TrainingConfig.save_optimizer_state and copied in at checkpoint export). -# -# Reuses the lotus_source_q50 run's query gradients and segment KFAC/lambda -# artifacts via symlinks; only the eigenfunction application (steps 6-7) and -# scoring (step 8) recompute. Same 6-checkpoint/3-segment recipe and lr_list. -# -# Validated LDS at N=400: mean Spearman rho=0.207 across the 50 queries -- -# weaker than the SGD variant's 0.390 on the same bank, suggesting the -# paper's diagonal-Hessian approximation costs more here than the -# optimizer-correct preconditioning gains. - -# Run with `bergson examples/method_comparison/gpt2_wikitext_lotus_source_adam_q50.yaml` - -steps: - - approxunrolling: - index_cfg: - run_path: runs/lotus_source_adam_q50 - # Tokenizer + final-model reference; per-checkpoint passes load the - # model dirs below. - model: runs/lotus/retrained/base - # In the approxunrolling pipeline `overwrite` doubles as resume: - # steps whose output dirs exist are skipped after a crash. - overwrite: true - precision: fp32 - # GPT-2's max context length is 1024. - token_batch_size: 1024 - # GPT-2 uses an nn.Linear for its LM head. - filter_modules: "lm_head" - distributed: - nproc_per_node: 8 # <-- adjust to this node's GPU count - nnode: 1 - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - # Doc-level scoring (one doc per row, no packing): packed 512-token - # chunks contaminate per-doc gradients across doc boundaries and - # were found to suppress LDS (see gpt2_lotus_trackstar50q_docs). - chunk_length: 0 - - hessian_cfg: - method: kfac - hessian_dtype: fp32 - ev_correction: true - - approx_unrolling_cfg: - checkpoints: - - runs/lotus_source_q50/models/step_51 - - runs/lotus_source_q50/models/step_102 - - runs/lotus_source_q50/models/step_136 - - runs/lotus_source_q50/models/step_187 - - runs/lotus_source_q50/models/step_238 - - runs/lotus_source_q50/models/step_288 - segments: 3 - lr_list: [4.857423e-04, 5.616667e-04, 2.416667e-04] - step_size_list: [97, 96, 96] - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[1:51]" - chunk_length: 0 - # One score column per query (50 total), like the EK-FAC and - # TrackStar runs, for multi-query LDS validation. - query_aggregation: none - # Adam variant: build per-segment diagonal preconditioners from the - # checkpoints' optimizer.pt files. - use_adam_preconditioner: true - - - validate: - run_path: runs/lotus_source_adam_q50_validate - model: gpt2 - overwrite: true - - scores: runs/lotus_source_adam_q50/scores.npy - retrained_dir: runs/lotus - - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 512 - - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[1:51]" - chunk_length: 0 - - batch_size: 64 diff --git a/examples/method_comparison/gpt2_wikitext_lotus_source_q50.yaml b/examples/method_comparison/gpt2_wikitext_lotus_source_q50.yaml deleted file mode 100644 index c7cafaba..00000000 --- a/examples/method_comparison/gpt2_wikitext_lotus_source_q50.yaml +++ /dev/null @@ -1,79 +0,0 @@ -# SOURCE (approximate unrolling, Bae et al. 2024) attribution scores for the -# lotus GPT-2 base model against the 50 test queries, validated against the -# complete 400-model lotus bank -- the setting where MAGIC scores rho=0.92. -# Plain SGD variant (no optimizer preconditioning); see -# gpt2_wikitext_lotus_source_adam_q50.yaml for the Adam/AdamW variant. -# -# Validated LDS at N=400: mean Spearman rho=0.390 across the 50 queries. - -# Run with `bergson examples/method_comparison/gpt2_wikitext_lotus_source_q50.yaml` - -steps: - - approxunrolling: - index_cfg: - run_path: runs/lotus_source_q50_damp0 - # Tokenizer + final-model reference; per-checkpoint passes load the - # model dirs below. - model: runs/lotus/retrained/base - # In the approxunrolling pipeline `overwrite` doubles as resume: - # steps whose output dirs exist are skipped after a crash. - overwrite: true - precision: fp32 - # GPT-2's max context length is 1024. - token_batch_size: 1024 - # GPT-2 uses an nn.Linear for its LM head. - filter_modules: "lm_head" - distributed: - nproc_per_node: 8 # <-- adjust to this node's GPU count - nnode: 1 - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - # Doc-level scoring (one doc per row, no packing): packed 512-token - # chunks contaminate per-doc gradients across doc boundaries and - # were found to suppress LDS (see gpt2_lotus_trackstar50q_docs). - chunk_length: 0 - - hessian_cfg: - method: kfac - hessian_dtype: fp32 - ev_correction: true - - approx_unrolling_cfg: - checkpoints: - - runs/lotus_source_q50/models/step_51 - - runs/lotus_source_q50/models/step_102 - - runs/lotus_source_q50/models/step_136 - - runs/lotus_source_q50/models/step_187 - - runs/lotus_source_q50/models/step_238 - - runs/lotus_source_q50/models/step_288 - segments: 3 - lr_list: [4.857423e-04, 5.616667e-04, 2.416667e-04] - step_size_list: [97, 96, 96] - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[1:51]" - chunk_length: 0 - # One score column per query (50 total), like the EK-FAC and - # TrackStar runs, for multi-query LDS validation. - query_aggregation: none - - - validate: - run_path: runs/lotus_source_q50_damp0_validate - model: gpt2 - overwrite: true - - scores: runs/lotus_source_q50_damp0/scores.npy - retrained_dir: runs/lotus - - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "train" - chunk_length: 512 - - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: "test[1:51]" - chunk_length: 0 - - batch_size: 64 diff --git a/examples/method_comparison/gpt2_wikitext_lotus_trackstar_q50.yaml b/examples/method_comparison/gpt2_wikitext_lotus_trackstar_q50.yaml deleted file mode 100644 index e71f91b5..00000000 --- a/examples/method_comparison/gpt2_wikitext_lotus_trackstar_q50.yaml +++ /dev/null @@ -1,170 +0,0 @@ -# TrackStar (Chang et al., 2024) attribution scores for the lotus GPT-2 -# base model against the 50 test queries, validated against the complete -# 400-model lotus bank -- the same setting as the MAGIC (rho=0.92) and -# SOURCE (rho=0.39) comparisons. Doc-level scoring, projection_dim 32, -# Adam-normalized gradients from runs/lotus/optimizer.pt, standard -# query/value hessian mixing. - -# Run with `bergson examples/method_comparison/gpt2_wikitext_lotus_trackstar_q50.yaml` - -steps: -- trackstar: - index_cfg: - run_path: runs/lotus_trackstar_q50 - overwrite: true - model: runs/lotus/retrained/base - precision: fp32 - revision: null - distributed: - nnode: 1 - nproc_per_node: 8 - node_rank: null - fsdp: false - peft_init_kwargs: '' - model_kwargs: '' - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: train - subset: null - prompt_column: text - completion_column: '' - conversation_column: '' - reward_column: '' - skip_nan_rewards: false - truncation: false - format_template: '' - data_kwargs: '' - chunk_length: 0 - tokenizer: '' - drop_columns: true - max_tokens: null - use_tf32_matmuls: false - debug: false - projection_dim: 32 - include_bias: false - reshape_to_square: false - projection_type: rademacher - projection_target: per_module - token_batch_size: 1024 - max_batch_size: null - auto_batch_size: false - optimizer_state: runs/lotus/optimizer.pt - loss_fn: ce - loss_reduction: sum - label_smoothing: 0.0 - stream_shard_size: 400000 - split_attention_modules: [] - attention: - num_heads: 0 - head_size: 0 - head_dim: 0 - profile: false - filter_modules: lm_head - force_math_sdp: false - attribute_tokens: false - modules: [] - trackstar_cfg: - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: test[1:51] - subset: null - prompt_column: text - completion_column: '' - conversation_column: '' - reward_column: '' - skip_nan_rewards: false - truncation: false - format_template: '' - data_kwargs: '' - chunk_length: 0 - preprocess_cfg: - unit_normalize: true - hessian_path: null - inversion_cfg: - inversion: damped_inverse - damping_factor: 0.1 - aggregation: none - normalize_aggregated_grad: false - score_cfg: - query_path: '' - score: individual - batch_size: 1024 - precision: fp32 - modules: [] - higher_is_better: true - target_downweight_components: 1000 - stats_sample_size: 10000 - resume: true -- validate: - run_path: runs/lotus_trackstar_q50_validate - overwrite: true - model: gpt2 - precision: fp32 - revision: null - distributed: - nnode: 1 - nproc_per_node: 8 - node_rank: null - fsdp: false - peft_init_kwargs: '' - model_kwargs: '' - data: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: train - subset: null - prompt_column: text - completion_column: '' - conversation_column: '' - reward_column: '' - skip_nan_rewards: false - truncation: false - format_template: '' - data_kwargs: '' - chunk_length: 512 - tokenizer: '' - drop_columns: true - max_tokens: null - use_tf32_matmuls: false - debug: false - lr_schedule: - lr: 1.0e-05 - lr_scheduler_type: linear - lr_start: 0.0 - lr_end: 0.0 - warmup_steps: 0 - num_cycles: 0.5 - power: 1.0 - batch_size: 64 - num_epochs: 1 - seed: 42 - adam_beta1: 0.95 - adam_beta2: 0.975 - eps_root: 1.0e-08 - optimizer: adamw - save_optimizer_state: false - weight_decay: 0.01 - max_grad_norm: null - grad_checkpointing: false - resume: false - wandb_project: '' - save_retrained_models: false - query: - dataset: EleutherAI/bergson-wikitext-512-chunks - split: test[1:51] - subset: null - prompt_column: text - completion_column: '' - conversation_column: '' - reward_column: '' - skip_nan_rewards: false - truncation: false - format_template: '' - data_kwargs: '' - chunk_length: 0 - query_method: mean - num_subsets: 100 - subset_strategy: random - exclude_zero_scores: false - subset_fraction: 0.0 - scores: runs/lotus_trackstar_q50/scores - retrained_dir: runs/lotus diff --git a/examples/pipelines/ekfac_qwen7b.yaml b/examples/pipelines/ekfac_qwen7b.yaml deleted file mode 100644 index e2a970ae..00000000 --- a/examples/pipelines/ekfac_qwen7b.yaml +++ /dev/null @@ -1,59 +0,0 @@ -# EKFAC influence pipeline for Qwen/Qwen2.5-7B-Instruct on a dummy test set. -# -# Fits an EK-FAC (KFAC + eigenvalue correction) Hessian over a small slice of -# pile-10k, then scores every index document against the mean gradient of a -# tiny 4-document query "test set". Small enough to sanity-check the pipeline -# end-to-end; scale up the dataset splits for a real run. -# -# Run with `bergson examples/pipelines/ekfac_qwen7b.yaml` -# -# Outputs land under runs/ekfac_qwen7b/. The score memmap is at -# runs/ekfac_qwen7b/scores/scores.bin; load it with bergson.data.load_scores. - -run_path: runs/ekfac_qwen7b -steps: - - ekfac: - index_cfg: - run_path: runs/ekfac_qwen7b - model: Qwen/Qwen2.5-7B-Instruct - overwrite: true - # bf16 keeps the 7B model comfortably within a single 48GB GPU. - precision: bf16 - token_batch_size: 2048 - # EK-FAC doesn't support gradient projection at any stage, so every - # gradient here is the full dense weight gradient. For a - # multi-billion-param model that means ONE document per batch — 16 - # docs would need hundreds of GB. Raise cautiously. - max_batch_size: 1 - # Qwen2.5's LM head is a huge (152k x 3584) nn.Linear; exclude it. - filter_modules: "lm_head" - distributed: - nproc_per_node: 8 - nnode: 1 - data: - # Dummy index corpus: a small slice of pile-10k. - dataset: NeelNanda/pile-10k - split: "train[:128]" - truncation: true - - hessian_cfg: - # KFAC + eigenvalue correction == EK-FAC. - method: kfac - ev_correction: true - - score_cfg: - batch_size: 512 - - preprocess_cfg: - unit_normalize: false - - hessian_pipeline_cfg: - query: - # Dummy query "test set": 8 held-out documents (>= nproc_per_node so - # every GPU is used; bergson caps workers to the dataset size). - dataset: NeelNanda/pile-10k - split: "train[128:136]" - truncation: true - query_aggregation: mean - inversion_cfg: - damping_factor: 0.1 diff --git a/examples/pipelines/recall_ekfac_olmo3.yaml b/examples/pipelines/recall_ekfac_olmo3.yaml deleted file mode 100644 index ff9aac96..00000000 --- a/examples/pipelines/recall_ekfac_olmo3.yaml +++ /dev/null @@ -1,63 +0,0 @@ -# EK-FAC recall for the LoRA-finetuned Olmo-3-7B on the synthetic -# factual-recall data, to compare against an existing TrackStar result. -# -# EK-FAC (kfac + ev_correction) scores every statement against every question -# (query_aggregation: none -> one score column per question), then `recall` -# ranks the gold statements per question and reports MRR / Recall@10. -# -# Scale is 100 people (3600 statements x 400 questions) to match the TrackStar -# run already on disk at: -# runs/recall_trackstar_olmo3_100p_1gpu/eval/summary.csv -# so the two are directly comparable. (TrackStar there: MRR 0.547, Recall@10 -# 0.812.) The model (a public LoRA adapter over allenai/Olmo-3-7B-Instruct) and -# its base are public, so no HF token is needed. -# -# Prerequisites: -# 1. A node with working multi-GPU (NCCL healthy): EK-FAC shards its KFAC -# factors across GPUs and OOMs on a single 48GB card for a 7B model. No -# `distributed:` block is set, so all GPUs on the node are used. -# 2. The 100-person recall datasets (already present on our filesystem; if not, -# regenerate them deterministically): -# python -m bergson.recall.generate --num_people 100 --seed 0 --data_dir data -# -# Run with `python -m bergson examples/pipelines/recall_ekfac_olmo3.yaml`, then: -# cat runs/recall_ekfac_olmo3_100p/eval/summary.csv # EK-FAC -# cat runs/recall_trackstar_olmo3_100p_1gpu/eval/summary.csv # TrackStar -# -# For a larger comparison, regenerate both datasets at a new --num_people M, -# rerun TrackStar at that M, and replace every `100` / `_100p_` below with M. - -run_path: runs/recall_ekfac_olmo3_100p -steps: - - ekfac: - index_cfg: - run_path: runs/recall_ekfac_olmo3_100p/ekfac - model: EleutherAI/Olmo-3-7B-Instruct-recall-9000 - precision: bf16 - overwrite: true - force_math_sdp: true - data: - dataset: data/statements_100p_seed0.hf - prompt_column: fact - hessian_cfg: - method: kfac - ev_correction: true - score_cfg: - batch_size: 1024 - preprocess_cfg: {} - hessian_pipeline_cfg: - query: - dataset: data/questions_100p_seed0.hf - prompt_column: text - query_aggregation: none - inversion_cfg: - inversion: damped_inverse - damping_factor: 0.1 - - - recall: - run_path: runs/recall_ekfac_olmo3_100p/eval - scores: runs/recall_ekfac_olmo3_100p/ekfac/scores - k: 10 - data: - num_people: 100 - seed: 0 diff --git a/examples/pipelines/recall_ekfac_olmo3_fp32.yaml b/examples/pipelines/recall_ekfac_olmo3_fp32.yaml deleted file mode 100644 index 890df2fd..00000000 --- a/examples/pipelines/recall_ekfac_olmo3_fp32.yaml +++ /dev/null @@ -1,62 +0,0 @@ -# EK-FAC recall for the LoRA-finetuned Olmo-3-7B on the synthetic -# factual-recall data, in fp32. -# -# Identical to recall_ekfac_olmo3.yaml except precision is fp32 (the model is -# loaded and gradients are collected in fp32 rather than bf16) and the run -# path has an _fp32 suffix. Compare against the fp32 TrackStar run at -# runs/recall_trackstar_olmo3_100p_fp32_1gpu/eval/summary.csv -# produced by recall_trackstar_olmo3_fp32.yaml. -# -# Prerequisites are the same as recall_ekfac_olmo3.yaml: a healthy multi-GPU -# node and the 100-person recall datasets: -# python -m bergson.recall.generate --num_people 100 --seed 0 --data_dir data -# -# Run with `python -m bergson examples/pipelines/recall_ekfac_olmo3_fp32.yaml`, -# then: -# cat runs/recall_ekfac_olmo3_100p_fp32/eval/summary.csv - -run_path: runs/recall_ekfac_olmo3_100p_fp32 -steps: - - ekfac: - index_cfg: - run_path: runs/recall_ekfac_olmo3_100p_fp32/ekfac - model: EleutherAI/Olmo-3-7B-Instruct-recall-9000 - precision: fp32 - overwrite: true - force_math_sdp: true - # Down from the 2048 default: in fp32 the 48GB A40s OOM during - # gradient collection at 2048, during the KFAC fit at 1024, and during - # scoring at 512 (per-sequence gradients for a batch are - # batch_size x 40M params fp32 next to the model + query slice). - token_batch_size: 256 - data: - dataset: data/statements_100p_seed0.hf - prompt_column: fact - hessian_cfg: - method: kfac - ev_correction: true - score_cfg: - batch_size: 1024 - # Score 64 queries per pass: the full 400-query set of unprojected - # fp32 gradients (~64GB) cannot sit on a 48GB card next to the model. - query_batch_size: 64 - preprocess_cfg: {} - hessian_pipeline_cfg: - # Skip pipeline steps whose output dir already exists (query gradients - # and KFAC factors survive a crash in a later step). - resume: true - query: - dataset: data/questions_100p_seed0.hf - prompt_column: text - query_aggregation: none - inversion_cfg: - inversion: damped_inverse - damping_factor: 0.1 - - - recall: - run_path: runs/recall_ekfac_olmo3_100p_fp32/eval - scores: runs/recall_ekfac_olmo3_100p_fp32/ekfac/scores - k: 10 - data: - num_people: 100 - seed: 0 diff --git a/examples/pipelines/recall_source_olmo3_fp32.yaml b/examples/pipelines/recall_source_olmo3_fp32.yaml deleted file mode 100644 index 49c94be1..00000000 --- a/examples/pipelines/recall_source_olmo3_fp32.yaml +++ /dev/null @@ -1,79 +0,0 @@ -# SOURCE (approximate unrolling, Bae et al.) recall for the LoRA-finetuned -# Olmo-3-7B on the synthetic factual-recall data, in fp32, to compare against -# the fp32 EK-FAC and TrackStar runs at the same 100-person scale: -# runs/recall_ekfac_olmo3_100p_fp32/eval/summary.csv -# runs/recall_trackstar_olmo3_100p_fp32_1gpu/eval/summary.csv -# -# 6 checkpoints / 3 segments over the 10125-step finetune -# (runs/finetune_recall_9000_olmo3: cosine LR peak 2e-4, 3% warmup, batch 128, -# 4 epochs). Checkpoints are HF adapter dirs exported from the trainer's DCP -# snapshots at steps {1700, 3400 | 5100, 6800 | 8400, 10124} — segment -# midpoints and ends for boundaries 0-3375-6750-10125: -# python scripts/export_dcp_checkpoints_to_hf.py \ -# --ckpt_dir runs/finetune_recall_9000_olmo3/checkpoints \ -# --base_model allenai/Olmo-3-7B-Instruct \ -# --peft_init_kwargs "r=16,lora_alpha=32,lora_dropout=0.0,target_modules=all-linear,task_type=CAUSAL_LM" \ -# --steps 1700 3400 5100 6800 8400 10124 \ -# --out_dir runs/recall_source_olmo3_100p_fp32/models -# -# lr_list is the per-third average of the training cosine schedule -# (LRScheduleConfig.get_schedule over 10125 steps); it cannot be inferred -# automatically because the MAGIC trainer writes no log_history.json. -# -# Disk: per-checkpoint covariances are ~66GB and segment eigenvectors ~66GB in -# fp32; expect a ~600GB peak. The per-ckpt covariance shards can be deleted -# once their segment's eigendecomposition (pipeline step 2/8) is done. -# -# Run with `python -m bergson examples/pipelines/recall_source_olmo3_fp32.yaml`, -# then: -# cat runs/recall_source_olmo3_100p_fp32/eval/summary.csv - -run_path: runs/recall_source_olmo3_100p_fp32 -steps: - - approxunrolling: - index_cfg: - run_path: runs/recall_source_olmo3_100p_fp32 - # Tokenizer + final-model reference; per-checkpoint passes load the - # adapter dirs below. - model: EleutherAI/Olmo-3-7B-Instruct-recall-9000 - # In the approxunrolling pipeline `overwrite` doubles as resume: - # steps whose output dirs exist are skipped after a crash. - overwrite: true - precision: fp32 - force_math_sdp: true - # fp32 memory ceiling on 48GB A40s, matching the EK-FAC fp32 run. - token_batch_size: 256 - data: - dataset: data/statements_100p_seed0.hf - prompt_column: fact - hessian_cfg: - method: kfac - hessian_dtype: fp32 - ev_correction: true - approx_unrolling_cfg: - checkpoints: - - runs/recall_source_olmo3_100p_fp32/models/step_1700 - - runs/recall_source_olmo3_100p_fp32/models/step_3400 - - runs/recall_source_olmo3_100p_fp32/models/step_5100 - - runs/recall_source_olmo3_100p_fp32/models/step_6800 - - runs/recall_source_olmo3_100p_fp32/models/step_8400 - - runs/recall_source_olmo3_100p_fp32/models/step_10124 - segments: 3 - lr_list: [1.770244e-04, 1.046429e-04, 1.833268e-05] - step_size_list: [3375, 3375, 3375] - query: - dataset: data/questions_100p_seed0.hf - prompt_column: text - # One score column per question, like the EK-FAC/TrackStar runs. - query_aggregation: none - # The 400 unaggregated fp32 query gradients (~60GB) cannot sit on a - # 48GB card next to the model during per-segment scoring. - query_batch_size: 64 - - - recall: - run_path: runs/recall_source_olmo3_100p_fp32/eval - scores: runs/recall_source_olmo3_100p_fp32/scores - k: 10 - data: - num_people: 100 - seed: 0 diff --git a/examples/pipelines/recall_trackstar_olmo3_fp32.yaml b/examples/pipelines/recall_trackstar_olmo3_fp32.yaml deleted file mode 100644 index 34821512..00000000 --- a/examples/pipelines/recall_trackstar_olmo3_fp32.yaml +++ /dev/null @@ -1,57 +0,0 @@ -# TrackStar recall for the LoRA-finetuned Olmo-3-7B on the synthetic -# factual-recall data, in fp32. -# -# Mirrors the config saved in runs/recall_trackstar_olmo3_100p_1gpu/config.yaml -# (bergson 0.10.0, git 989808d0) with precision fp32 instead of bf16 and an -# _fp32 run path. Kept single-process (nproc_per_node: 1) because the cosine -# recall variants OOM on host RAM when run multi-GPU. Compare against the fp32 -# EK-FAC run from recall_ekfac_olmo3_fp32.yaml. -# -# Run with `python -m bergson examples/pipelines/recall_trackstar_olmo3_fp32.yaml`, -# then: -# cat runs/recall_trackstar_olmo3_100p_fp32_1gpu/eval/summary.csv - -run_path: runs/recall_trackstar_olmo3_100p_fp32_1gpu -steps: - - trackstar: - index_cfg: - run_path: runs/recall_trackstar_olmo3_100p_fp32_1gpu/trackstar - overwrite: true - model: EleutherAI/Olmo-3-7B-Instruct-recall-9000 - precision: fp32 - distributed: - nnode: 1 - nproc_per_node: 1 - data: - dataset: data/statements_100p_seed0.hf - prompt_column: fact - projection_dim: 16 - projection_type: rademacher - projection_target: per_module - # Down from 2048 in the bf16 run: the fp32 model + backward OOMs a - # single 48GB A40 at that token budget. - token_batch_size: 512 - force_math_sdp: true - trackstar_cfg: - query: - dataset: data/questions_100p_seed0.hf - prompt_column: text - preprocess_cfg: - unit_normalize: true - inversion_cfg: - inversion: damped_inverse - damping_factor: 0.1 - aggregation: none - score_cfg: - score: individual - batch_size: 1024 - precision: fp32 - target_downweight_components: 1000 - - - recall: - run_path: runs/recall_trackstar_olmo3_100p_fp32_1gpu/eval - scores: runs/recall_trackstar_olmo3_100p_fp32_1gpu/trackstar/scores - k: 10 - data: - num_people: 100 - seed: 0 diff --git a/examples/pipelines/recall_trackstar_olmo3_fp32_dim32.yaml b/examples/pipelines/recall_trackstar_olmo3_fp32_dim32.yaml deleted file mode 100644 index 2cf661be..00000000 --- a/examples/pipelines/recall_trackstar_olmo3_fp32_dim32.yaml +++ /dev/null @@ -1,57 +0,0 @@ -# TrackStar recall in fp32 with projection_dim 32 (vs 16 in -# recall_trackstar_olmo3_fp32.yaml, otherwise identical). Each module gradient -# gets a double-sided 32x32 random projection instead of 16x16, i.e. 4x the -# projected feature dims — testing whether TrackStar's gap to EK-FAC/SOURCE at -# dim 16 (see experiments/recall_attribution_fp32/README.md) is a projection -# bottleneck. dim 64 OOMs: the dense per-module Grams are (d^2)^2 fp32, ~30GB -# at d=64 next to the 28GB model. d=32 keeps them at ~1.9GB. -# -# Run with -# `python -m bergson examples/pipelines/recall_trackstar_olmo3_fp32_dim32.yaml`, -# then: -# cat runs/recall_trackstar_olmo3_100p_fp32_dim32/eval/summary.csv - -run_path: runs/recall_trackstar_olmo3_100p_fp32_dim32 -steps: - - trackstar: - index_cfg: - run_path: runs/recall_trackstar_olmo3_100p_fp32_dim32/trackstar - overwrite: true - model: EleutherAI/Olmo-3-7B-Instruct-recall-9000 - precision: fp32 - distributed: - nnode: 1 - nproc_per_node: 8 - data: - dataset: data/statements_100p_seed0.hf - prompt_column: fact - projection_dim: 32 - projection_type: rademacher - projection_target: per_module - # Down from 2048 in the bf16 run: the fp32 model + backward OOMs a - # single 48GB A40 at that token budget. - token_batch_size: 512 - force_math_sdp: true - trackstar_cfg: - query: - dataset: data/questions_100p_seed0.hf - prompt_column: text - preprocess_cfg: - unit_normalize: true - inversion_cfg: - inversion: damped_inverse - damping_factor: 0.1 - aggregation: none - score_cfg: - score: individual - batch_size: 1024 - precision: fp32 - target_downweight_components: 1000 - - - recall: - run_path: runs/recall_trackstar_olmo3_100p_fp32_dim32/eval - scores: runs/recall_trackstar_olmo3_100p_fp32_dim32/trackstar/scores - k: 10 - data: - num_people: 100 - seed: 0 diff --git a/tests/test_optimizer_placement.py b/tests/test_optimizer_placement.py deleted file mode 100644 index 373c7186..00000000 --- a/tests/test_optimizer_placement.py +++ /dev/null @@ -1,205 +0,0 @@ -"""Placing trainer-written optimizer states where SOURCE and TrackStar read them. - -The trainer writes ``step_.optimizer.pt`` as a sibling of ``step_.ckpt``; -both consumers want ``optimizer.pt`` *inside* a directory. These tests pin the -matching rules, since a wrong match silently builds a preconditioner from the -wrong step rather than failing. -""" - -import os -import shutil -from pathlib import Path - -import pytest -import torch - -from bergson.approx_unrolling.adam_preconditioner import OPTIMIZER_STATE_FILE -from bergson.utils.load_from_optimizer import load_optimizer -from bergson.utils.optimizer_placement import ( - place_final_optimizer_state, - place_optimizer_states, - sorted_optimizer_states, -) - - -def _write_state(path: Path, step: int) -> None: - """A minimal optimizer.pt whose payload identifies its step.""" - path.parent.mkdir(parents=True, exist_ok=True) - torch.save( - { - "state": {0: {"exp_avg_sq": torch.full((2, 3), float(step))}}, - "param_groups": [{"betas": (0.9, 0.999), "eps": 1e-8, "lr": 1e-4}], - }, - path, - ) - - -@pytest.fixture -def save_dir(tmp_path) -> Path: - """A trainer output dir with sibling states at steps 0, 4 and 9.""" - d = tmp_path / "run" - for step in (0, 4, 9): - _write_state(d / f"step_{step}.optimizer.pt", step) - (d / f"step_{step}.ckpt").mkdir(parents=True, exist_ok=True) - return d - - -def _exported(root: Path, names) -> list[Path]: - dirs = [] - for name in names: - p = root / name - p.mkdir(parents=True, exist_ok=True) - dirs.append(p) - return dirs - - -def test_sorted_optimizer_states_orders_numerically(save_dir): - """Steps must sort numerically, not lexically (step_10 < step_9 as strings).""" - _write_state(save_dir / "step_10.optimizer.pt", 10) - assert [s for s, _ in sorted_optimizer_states(save_dir)] == [0, 4, 9, 10] - - -def test_sorted_optimizer_states_ignores_other_files(save_dir): - """The final optimizer.pt and the .ckpt dirs are not per-step states.""" - _write_state(save_dir / OPTIMIZER_STATE_FILE, -1) - assert [s for s, _ in sorted_optimizer_states(save_dir)] == [0, 4, 9] - - -def test_places_by_step_name(save_dir, tmp_path): - """Directories naming a step get that step's state, regardless of order.""" - dsts = _exported(tmp_path / "hf", ["step_9", "step_0", "step_4"]) - placed = place_optimizer_states(save_dir, dsts, mode="copy") - - assert len(placed) == 3 - for dst, expected in zip(dsts, (9, 0, 4)): - blob = load_optimizer(str(dst)) - assert blob["state"][0]["exp_avg_sq"][0, 0].item() == expected - - -def test_places_positionally_when_unnamed(save_dir, tmp_path): - """Unnamed dirs are matched in step order, so seg 0 gets the earliest step.""" - dsts = _exported(tmp_path / "hf", ["a", "b", "c"]) - place_optimizer_states(save_dir, dsts, mode="copy") - - for dst, expected in zip(dsts, (0, 4, 9)): - blob = load_optimizer(str(dst)) - assert blob["state"][0]["exp_avg_sq"][0, 0].item() == expected - - -def test_positional_requires_equal_counts(save_dir, tmp_path): - """A short/long dir list means the caller passed the wrong run; don't guess.""" - dsts = _exported(tmp_path / "hf", ["a", "b"]) - with pytest.raises(ValueError, match="Cannot match positionally"): - place_optimizer_states(save_dir, dsts, mode="copy") - - -def test_mixed_naming_is_rejected(save_dir, tmp_path): - """Half-named lists can't be matched reliably, so refuse rather than guess.""" - dsts = _exported(tmp_path / "hf", ["step_0", "b", "c"]) - with pytest.raises(ValueError, match="name a step and some do not"): - place_optimizer_states(save_dir, dsts, mode="copy") - - -def test_unknown_step_is_rejected(save_dir, tmp_path): - """A named step with no matching state is an error, not a skip.""" - dsts = _exported(tmp_path / "hf", ["step_0", "step_7"]) - with pytest.raises(FileNotFoundError, match=r"steps \[7\]"): - place_optimizer_states(save_dir, dsts, mode="copy") - - -def test_refuses_to_clobber_without_overwrite(save_dir, tmp_path): - """Existing optimizer.pt files are protected; nothing is written on refusal.""" - dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) - (dsts[1] / OPTIMIZER_STATE_FILE).write_text("do not clobber") - - with pytest.raises(FileExistsError): - place_optimizer_states(save_dir, dsts, mode="copy") - - assert (dsts[1] / OPTIMIZER_STATE_FILE).read_text() == "do not clobber" - # The check precedes all writes, so the untouched dirs stay untouched. - assert not (dsts[0] / OPTIMIZER_STATE_FILE).exists() - - place_optimizer_states(save_dir, dsts, mode="copy", overwrite=True) - assert load_optimizer(str(dsts[1]))["state"][0]["exp_avg_sq"][0, 0].item() == 4 - - -def test_symlink_mode_is_relative_and_loadable(save_dir, tmp_path): - """Default mode links rather than duplicating a state per checkpoint.""" - dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) - place_optimizer_states(save_dir, dsts) - - link = dsts[0] / OPTIMIZER_STATE_FILE - assert link.is_symlink() - assert not os.path.isabs(os.readlink(link)), "link must be relative" - assert load_optimizer(str(dsts[0]))["state"][0]["exp_avg_sq"][0, 0].item() == 0 - - -def test_symlinks_survive_relocation(save_dir, tmp_path): - """Relative links keep resolving when run and checkpoints move together.""" - dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) - place_optimizer_states(save_dir, dsts) - - moved = tmp_path / "moved" - moved.mkdir() - shutil.copytree(tmp_path / "hf", moved / "hf", symlinks=True) - shutil.copytree(save_dir, moved / save_dir.name, symlinks=True) - - relocated = moved / "hf" / "step_0" / OPTIMIZER_STATE_FILE - assert relocated.is_symlink() - assert relocated.resolve().is_file(), "relative link broke after relocation" - assert load_optimizer(str(moved / "hf" / "step_0"))["state"][0][ - "exp_avg_sq" - ][0, 0].item() == 0 - - -def test_hardlink_and_move_modes(save_dir, tmp_path): - dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) - place_optimizer_states(save_dir, dsts, mode="hardlink") - assert load_optimizer(str(dsts[2]))["state"][0]["exp_avg_sq"][0, 0].item() == 9 - - dsts2 = _exported(tmp_path / "hf2", ["step_0", "step_4", "step_9"]) - place_optimizer_states(save_dir, dsts2, mode="move") - assert not (save_dir / "step_9.optimizer.pt").exists() - assert load_optimizer(str(dsts2[2]))["state"][0]["exp_avg_sq"][0, 0].item() == 9 - - -def test_missing_states_names_the_config_flag(tmp_path): - """The error should say how to produce the files, not just that they're absent.""" - empty = tmp_path / "empty" - empty.mkdir() - dsts = _exported(tmp_path / "hf", ["step_0"]) - with pytest.raises(FileNotFoundError, match="save_optimizer_state"): - place_optimizer_states(empty, dsts) - - -def test_missing_destination_is_rejected(save_dir, tmp_path): - with pytest.raises(NotADirectoryError): - place_optimizer_states(save_dir, [tmp_path / "nope"]) - - -def test_place_final_state_for_trackstar(save_dir, tmp_path): - """TrackStar points optimizer_state at a dir; the final state must land there.""" - _write_state(save_dir / OPTIMIZER_STATE_FILE, 99) - exported = _exported(tmp_path / "model", ["final"])[0] - - target = place_final_optimizer_state(save_dir, exported, mode="copy") - - assert target == exported / OPTIMIZER_STATE_FILE - # load_optimizer resolves a directory, which is what PreprocessConfig takes. - assert load_optimizer(str(exported))["state"][0]["exp_avg_sq"][0, 0].item() == 99 - - -def test_place_final_state_is_idempotent_in_place(save_dir): - """Pointing it at its own directory is a no-op, not a self-clobber.""" - _write_state(save_dir / OPTIMIZER_STATE_FILE, 99) - target = place_final_optimizer_state(save_dir, save_dir) - assert target.is_file() - assert load_optimizer(str(save_dir))["state"][0]["exp_avg_sq"][0, 0].item() == 99 - - -def test_place_final_state_missing_names_the_flag(tmp_path): - empty = tmp_path / "empty" - empty.mkdir() - dst = _exported(tmp_path / "model", ["final"])[0] - with pytest.raises(FileNotFoundError, match="save_optimizer_state"): - place_final_optimizer_state(empty, dst) From 45ff292014b115a0dbdb5b54ba4d0e8fe90da160 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 30 Jul 2026 07:28:25 +0000 Subject: [PATCH 09/10] refactor: drop the save_retrained_models alias The rename is done and every config in the repo uses save_models, so the deprecated alias, its conflict check and the FutureWarning are just weight. Also trims the field's docstring to the three destinations it actually writes. --- bergson/config/config.py | 37 +++---------------------------------- 1 file changed, 3 insertions(+), 34 deletions(-) diff --git a/bergson/config/config.py b/bergson/config/config.py index a1a1925e..0730f480 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -1,6 +1,5 @@ import math import os -import warnings from abc import ABC from dataclasses import dataclass from pathlib import Path @@ -381,24 +380,9 @@ class TrainingConfig(AttributionConfig, Serializable): """Weights & Biases project name. If set, logs training loss to W&B.""" save_models: bool = False - """When True, persist the models this run trains, in the form each run - produces them: - - - Training and validation runs save the fully-trained model (HF format, - weights + tokenizer) to ``/retrained/base/``, and validation - additionally saves each leave-k-out model to - ``/retrained/subset_/`` so it can be reused for later - attribution queries without retraining. ~0.5 GB per model for GPT-2. - - Metasmoothness runs save the unperturbed run's parameters to - ``/theta_final.pt`` and its initialization to ``theta_init.pt`` - (name -> fp32 CPU tensor, ~0.65 GB each for GPT-2). The two perturbed - runs are finite-difference probes and are not saved.""" - - # TODO(Lucia Quirke): remove by 2026-10-29. Backwards compatibility for the - # pre-rename key; the flag moved from ValidationConfig to TrainingConfig, - # where "retrained" no longer describes a plain training run. - save_retrained_models: bool | None = None - """Deprecated alias for ``save_models``.""" + """Persist the models this run trains: HF format under + ``/retrained/`` for training and validation, ``theta_init.pt`` and + ``theta_final.pt`` for metasmoothness. ~0.5 GB per model for GPT-2.""" def __post_init__(self): super().__post_init__() @@ -407,21 +391,6 @@ def __post_init__(self): if isinstance(self.save_optimizer_state, bool): self.save_optimizer_state = "last" if self.save_optimizer_state else "none" - # TODO(Lucia Quirke): remove by 2026-10-29, along with the field above. - if self.save_retrained_models is not None: - if self.save_models and self.save_models != self.save_retrained_models: - raise ValueError( - "save_retrained_models and save_models are set to conflicting " - "values; save_retrained_models is a deprecated alias, so set " - "only save_models." - ) - warnings.warn( - "save_retrained_models is deprecated; use save_models.", - FutureWarning, - stacklevel=2, - ) - self.save_models = self.save_retrained_models - @dataclass class MetasmoothnessConfig(TrainingConfig): From aae67d5f0d9da49ec7b1c285a0bc7895f00153cf Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Sat, 1 Aug 2026 08:46:39 +0900 Subject: [PATCH 10/10] Add YAMLs for Adam WikiText SOURCE + EK-FAC replications (#369) * feat: Fisher normalization for SOURCE segment eigenvalues Segment eigenvalues were unnormalized sums over (documents x checkpoints). EK-FAC is unaffected -- every inversion is homogeneous of degree -1 in lambda because damping is relative to mean(lambda) -- but SOURCE's eigenfunctions are not scale-invariant, and sigma enters exp(-lr*K*sigma). Eq. 1 of Bae et al. 2024 defines the empirical risk as a mean over data points, so fisher_normalization defaults to per-document; token units put the mask argument in its dead-linear range. Also makes precompute_checkpoints resumable rather than overwrite-only, and lets resolve() fall back when a checkpoint's sibling config.yaml turns out to belong to an attribution run. * feat: per-snapshot optimizer.pt export and interval save mode in MAGIC trainer (cherry picked from commit 16c29730ff6fdadd1c61ea72dc9dda43e6a85274) * fix: restrict the end-of-training snapshot to interval save mode The other save modes' backward replay indexes the data stream by checkpoint, so an extra trailing snapshot runs it past the end (DataStream index out of range in test_fsdp_ddp_scores_match). * fix: make skip_metagradient actually skip the MAGIC backward The flag only toggled allow_compile; the backward still ran (and crashed on interval-mode end-of-training snapshots). Port the dummy-zero-scores skip path so bank-only builds go straight to validation. (cherry picked from commit 13b631d3b030be312a881ceff6c31475cbceb4b5) * feat: skip_metagradient config field for trainer-only MAGIC runs Companion to the cherry-picked gate: defines the MagicConfig field and the do_metagradient condition it reads. * feat: Eq-43 hybrid for the Adam SOURCE segment eigenfunction Bae et al. 2024 App. D: 'use the diagonal Hessian approximation for computing the matrix exponential in Equation 39 and Equation 43. Note that we still use the EK-FAC factors to compute H^-1 g in Equation 43.' Evaluating the eigenfunction wholly on the diagonal drops that second half. With adam_segment_hybrid the diagonal supplies only the exponential mask (f_one_minus_exp) and a chained EK-FAC inverse supplies 1/x, applied to the query as the adjoint H^-1 @ M_mask @ q. * docs: WikiText-2 replication configs (train / EK-FAC / SOURCE) Paper recipe from Bae et al. 2024 App. B.1 (AdamW, lr 3e-5, wd 1e-2, batch 8, seq 512, 6 checkpoints) with train_mode enabled so GPT-2's 0.1 dropout is active, matching kronfluence's train.py. SOURCE uses the App. C preconditioned variant with the App. D Eq-43 hybrid, L=3, per-document Fisher units. * feat: LDS against a precomputed subset-retrain ground truth Scoring side of a published retrain bank (kronfluence ships masks.pt / losses.pt for its WikiText-2 experiment): sum each subset's member scores, Spearman against the measured subset losses, average over queries. Lets a replication be checked against published numbers without retraining. * fix: read the training config out of the saved {steps, metadata} document save_run_config has always wrapped the step list in a document with a metadata block, so load_training_config rejected every real run directory; its test hand-rolled the bare list it expected instead of calling save_run_config. Have the test write configs the way the writer does, and report an unparseable payload as ValueError so callers guessing at a run dir still degrade to their fallback. * docs: driver for the WikiText-2 replication chain Export -> EK-FAC IF -> SOURCE -> LDS against kronfluence's shipped bank. EK-FAC 0.4451 against their reported 0.44; SOURCE 0.4506, above IF as in the paper's Figure 6. * style: apply black formatting --- .../approx_unrolling/approx_unrolling_math.py | 137 ++++++++++++------ bergson/approx_unrolling/pipeline.py | 19 ++- .../precompute_checkpoints.py | 65 ++++++--- .../approx_unrolling/segment_aggregation.py | 71 ++++++++- bergson/approx_unrolling/trainer_run.py | 31 +++- bergson/config/config.py | 20 +++ bergson/hessians/apply_hessian.py | 32 +++- bergson/magic/cli.py | 21 ++- bergson/magic/config.py | 15 +- bergson/magic/metasmoothness.py | 87 +++++++++++ bergson/magic/trainer.py | 29 +++- bergson/validate.py | 56 +++++++ ...pt2_wikitext_lotus_source_q50_docnorm.yaml | 83 +++++++++++ ...ext_lotus_source_q50_docnorm_validate.yaml | 18 +++ ...pt2_wikitext_lotus_source_q50_toknorm.yaml | 83 +++++++++++ examples/replication/wikitext_gpt2_ekfac.yaml | 56 +++++++ .../replication/wikitext_gpt2_source.yaml | 76 ++++++++++ examples/replication/wikitext_gpt2_train.yaml | 67 +++++++++ .../chain_toknorm.sh | 50 +++++++ .../compare_lds.py | 81 +++++++++++ .../source_fisher_normalization/setup_run.py | 107 ++++++++++++++ experiments/wikitext_replication/run_chain.py | 64 ++++++++ tests/test_lds_precomputed_subsets.py | 78 ++++++++++ tests/test_source_fisher_normalization.py | 83 +++++++++++ tests/test_source_trainer_integration.py | 44 +++++- 25 files changed, 1388 insertions(+), 85 deletions(-) create mode 100644 examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm.yaml create mode 100644 examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm_validate.yaml create mode 100644 examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml create mode 100644 examples/replication/wikitext_gpt2_ekfac.yaml create mode 100644 examples/replication/wikitext_gpt2_source.yaml create mode 100644 examples/replication/wikitext_gpt2_train.yaml create mode 100755 experiments/source_fisher_normalization/chain_toknorm.sh create mode 100644 experiments/source_fisher_normalization/compare_lds.py create mode 100644 experiments/source_fisher_normalization/setup_run.py create mode 100644 experiments/wikitext_replication/run_chain.py create mode 100644 tests/test_lds_precomputed_subsets.py create mode 100644 tests/test_source_fisher_normalization.py diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index 4ec7b518..ddc886fc 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -3,7 +3,7 @@ import shutil from copy import deepcopy from pathlib import Path -from typing import Callable +from typing import Callable, Literal import numpy as np import torch @@ -129,21 +129,37 @@ def fn(sigma: Tensor) -> Tensor: return fn +def f_one_minus_exp(lr_times_steps: float) -> Callable[[Tensor], Tensor]: + """x -> 1 - exp(-lr_times_steps*x), the numerator of :func:`f_segment`. + + Used by the Eq-43 hybrid: the 1/x that completes f_segment is supplied by + the chained EK-FAC inverse rather than evaluated on the diagonal. + """ + + def fn(sigma: Tensor) -> Tensor: + return -torch.expm1(-lr_times_steps * sigma) + + return fn + + def apply_eigfn_to_query( src_grad_path: Path, dst_grad_path: Path, segment_dir: Path, lr_times_steps: float, - fn_kind: str, + fn_kind: Literal["f_segment", "f_backward"], distributed: DistributedConfig, preconditioner_path: str = "", + preconditioner_hybrid: bool = False, + preconditioner_hybrid_damping: float = 0.1, ) -> None: - """Apply F_segment or F_backward of one segment to a stored query gradient. + """Apply a segment's f_segment or f_backward operator to a stored query + gradient. The operator scales each query element in the segment's + eigenbasis by the function's value at the matching eigenvalue, averaged + over the segment's (checkpoint, document) pairs by + ``ApproxUnrollingConfig.fisher_normalization``. - ``fn_kind`` is "f_segment" or "f_backward". The segment eigenvalues are - already checkpoint-averaged (expected eigenvalues), so the eigenfunction is - applied to them directly. - ``preconditioner_path`` selects the preconditioned-optimizer variant: the + When ``preconditioner_path`` is set for an Adam-optimized run the eigenfunction is evaluated on a diagonal approximation of P^1/2 H P^1/2 in parameter space, with F_segment's output additionally multiplied by P (its P^1/2 . P^1/2 sandwich; F_backward's sandwich cancels).""" @@ -153,7 +169,11 @@ def apply_eigfn_to_query( run_path=str(dst_grad_path), ev_correction=True, preconditioner_path=preconditioner_path, - preconditioner_post_multiply=fn_kind == "f_segment", + # The hybrid's P^1/2 sandwich is absorbed by the chained EK-FAC inverse. + preconditioner_post_multiply=fn_kind == "f_segment" + and not preconditioner_hybrid, + preconditioner_hybrid=preconditioner_hybrid and fn_kind == "f_segment", + preconditioner_hybrid_damping=preconditioner_hybrid_damping, ) launch_distributed_run( "apply_eigfn_to_query", @@ -169,13 +189,17 @@ def _apply_eigfn_worker( world_size: int, cfg: EkfacConfig, lr_times_steps: float, - fn_kind: str, + fn_kind: Literal["f_segment", "f_backward"], ) -> None: init_dist(rank, local_rank, world_size) # Segment eigenvalues are already checkpoint-averaged, so the eigenfunction # is applied to them directly (no per-example normalization). - fn = {"f_segment": f_segment, "f_backward": f_backward}[fn_kind](lr_times_steps) + if cfg.preconditioner_hybrid: + # 1/x comes from the chained EK-FAC inverse, so mask with the numerator. + fn = f_one_minus_exp(lr_times_steps) + else: + fn = {"f_segment": f_segment, "f_backward": f_backward}[fn_kind](lr_times_steps) EkfacApplicator(cfg, apply_fn=fn).compute_ivhp_sharded() @@ -185,6 +209,8 @@ def walk_query_phase1( lr_times_steps_per_segment: list[float], distributed: DistributedConfig, preconditioner_paths: list[str] | None = None, + preconditioner_hybrid: bool = False, + preconditioner_hybrid_damping: float = 0.1, ) -> list[Path]: """Phase 1: build query_grad_0, ..., query_grad_{L-1} by walking F_backward. @@ -211,6 +237,8 @@ def walk_query_phase1( fn_kind="f_backward", distributed=distributed, preconditioner_path=preconditioner_paths[k] if preconditioner_paths else "", + preconditioner_hybrid=preconditioner_hybrid, + preconditioner_hybrid_damping=preconditioner_hybrid_damping, ) query_grad_paths[k - 1] = dst @@ -224,6 +252,8 @@ def walk_query_phase2( query_grad_paths: list[Path], distributed: DistributedConfig, preconditioner_paths: list[str] | None = None, + preconditioner_hybrid: bool = False, + preconditioner_hybrid_damping: float = 0.1, ) -> list[Path]: """Phase 2: build query_grad_segment_0, ..., query_grad_segment_{L-1} via F_segment. @@ -248,6 +278,8 @@ def walk_query_phase2( fn_kind="f_segment", distributed=distributed, preconditioner_path=preconditioner_paths[l] if preconditioner_paths else "", + preconditioner_hybrid=preconditioner_hybrid, + preconditioner_hybrid_damping=preconditioner_hybrid_damping, ) query_grad_segment_paths.append(dst) @@ -257,51 +289,64 @@ def walk_query_phase2( def score_per_segment_and_aggregate( index_cfg: IndexConfig, query_grad_segment_paths: list[Path], - final_checkpoint: str, + segment_checkpoints: list[list[str]], query_batch_size: int | None = None, ) -> Path: - """Phase 3: per-segment ``query_grad_segment_l . g(z_m)`` scores, summed. - - For each l, runs :func:`score_dataset` against the training data at the - final checkpoint with ``query_grad_segment_l`` as the query. Writes - per-segment outputs to ``/segment_{l}/scores/``, then sums into - ``/scores/``. + """Phase 3: per-segment ``query_grad_segment_l . g_bar_l(z_m)`` scores, summed. + + ``g_bar_l`` is the segment's expected training gradient (Bae et al. 2024, + Sec. 3.4; App. D notes the scores need "the training gradients C times, + where C is the total number of checkpoints"). Per-example gradients are + computed at each checkpoint *within* segment l and averaged; scoring every + segment at the final checkpoint instead collapses SOURCE toward influence + functions at convergence. + + Scores are linear in the training gradient, so averaging the per-checkpoint + scores is equivalent to scoring against the averaged gradient. Per-checkpoint + outputs land at ``/segment_{l}/scores_ckpt_{c}/``. """ base_run = Path(index_cfg.run_path) - num_segments = len(query_grad_segment_paths) score_dirs: list[Path] = [] - for l in range(num_segments): - scores_dir = base_run / f"segment_{l}" / "scores" - if index_cfg.distributed._node_rank == 0 and scores_dir.exists(): - shutil.rmtree(scores_dir) - seg_index_cfg = deepcopy(index_cfg) - seg_index_cfg.model = final_checkpoint - seg_index_cfg.run_path = str(scores_dir) - seg_index_cfg.projection_dim = 0 - score_cfg = ScoreConfig( - query_path=str(query_grad_segment_paths[l]), - higher_is_better=True, - query_batch_size=query_batch_size, - ) - seg_preprocess_cfg = PreprocessConfig() - save_run_config( - Score(score_cfg, seg_index_cfg, seg_preprocess_cfg), - seg_index_cfg.partial_run_path, - ) - score_dataset(seg_index_cfg, score_cfg, seg_preprocess_cfg) - score_dirs.append(scores_dir) - total = None - for scores_dir in score_dirs: - seg_scores = load_scores(scores_dir)[:] - seg_score_cfg = load_subconfig(scores_dir, "score_cfg", ScoreConfig) - if seg_score_cfg is None: + def _oriented(scores_dir: Path): + scores = load_scores(scores_dir)[:] + score_cfg = load_subconfig(scores_dir, "score_cfg", ScoreConfig) + if score_cfg is None: raise FileNotFoundError( f"No score_cfg found at {scores_dir}; cannot determine the " - "segment scores' orientation for aggregation." + "scores' orientation for aggregation." + ) + return -scores if score_cfg.higher_is_better else scores + + total = None + for l, ckpts in enumerate(segment_checkpoints): + seg_total = None + for c, ckpt in enumerate(ckpts): + scores_dir = base_run / f"segment_{l}" / f"scores_ckpt_{c}" + if index_cfg.distributed._node_rank == 0 and scores_dir.exists(): + shutil.rmtree(scores_dir) + seg_index_cfg = deepcopy(index_cfg) + seg_index_cfg.model = ckpt + seg_index_cfg.run_path = str(scores_dir) + seg_index_cfg.projection_dim = 0 + score_cfg = ScoreConfig( + query_path=str(query_grad_segment_paths[l]), + higher_is_better=True, + query_batch_size=query_batch_size, + ) + seg_preprocess_cfg = PreprocessConfig() + save_run_config( + Score(score_cfg, seg_index_cfg, seg_preprocess_cfg), + seg_index_cfg.partial_run_path, ) - if seg_score_cfg.higher_is_better: - seg_scores = -seg_scores + score_dataset(seg_index_cfg, score_cfg, seg_preprocess_cfg) + score_dirs.append(scores_dir) + + ckpt_scores = _oriented(scores_dir) + seg_total = ckpt_scores if seg_total is None else seg_total + ckpt_scores + + assert seg_total is not None, "each segment has >= 1 checkpoint" + seg_scores = seg_total / len(ckpts) total = seg_scores if total is None else total + seg_scores assert total is not None, "num_segments >= 1 is validated by the pipeline" diff --git a/bergson/approx_unrolling/pipeline.py b/bergson/approx_unrolling/pipeline.py index 971ea337..88f8e872 100644 --- a/bergson/approx_unrolling/pipeline.py +++ b/bergson/approx_unrolling/pipeline.py @@ -141,7 +141,7 @@ def approx_unrolling_pipeline( index_cfg, hessian_cfg, approx_unrolling_cfg, - overwrite=index_cfg.overwrite, + resume=index_cfg.overwrite, ) # Encourage GC between expensive steps; matters most in single-GPU mode # where the worker ran in-process and may still hold model references. @@ -185,6 +185,7 @@ def approx_unrolling_pipeline( per_segment=n_ckpts // n_segments, distributed=index_cfg.distributed, resume=index_cfg.overwrite, + normalization=approx_unrolling_cfg.fisher_normalization, ) # ── Step 5: Mean query gradient at the final checkpoint @@ -252,6 +253,8 @@ def approx_unrolling_pipeline( lr_times_steps_per_segment=lr_times_steps_per_segment, distributed=index_cfg.distributed, preconditioner_paths=[str(p) for p in preconditioner_paths] or None, + preconditioner_hybrid=approx_unrolling_cfg.adam_segment_hybrid, + preconditioner_hybrid_damping=approx_unrolling_cfg.adam_hybrid_damping, ) # ── Step 7: Phase 2 -- Get per-ckpt queries from segment queries @@ -267,6 +270,8 @@ def approx_unrolling_pipeline( query_grad_paths=query_grad_paths, distributed=index_cfg.distributed, preconditioner_paths=[str(p) for p in preconditioner_paths] or None, + preconditioner_hybrid=approx_unrolling_cfg.adam_segment_hybrid, + preconditioner_hybrid_damping=approx_unrolling_cfg.adam_hybrid_damping, ) # ── Step 8: Phase 3 -- per-segment scoring + sum @@ -276,7 +281,17 @@ def approx_unrolling_pipeline( out_path = score_per_segment_and_aggregate( index_cfg=index_cfg, query_grad_segment_paths=query_grad_segment_paths, - final_checkpoint=str(approx_unrolling_cfg.checkpoints[-1]), + # Bae et al. Sec. 3.4: g_bar_l averages the training gradients over the + # checkpoints within segment l, not the final checkpoint alone. + segment_checkpoints=[ + [ + str(c) + for c in approx_unrolling_cfg.checkpoints[ + l * (n_ckpts // n_segments) : (l + 1) * (n_ckpts // n_segments) + ] + ] + for l in range(n_segments) + ], query_batch_size=approx_unrolling_cfg.query_batch_size, ) logger.info(f"[approximate unrolling pipeline] DONE. Final scores at {out_path}") diff --git a/bergson/approx_unrolling/precompute_checkpoints.py b/bergson/approx_unrolling/precompute_checkpoints.py index a61ef7a3..8db2d579 100644 --- a/bergson/approx_unrolling/precompute_checkpoints.py +++ b/bergson/approx_unrolling/precompute_checkpoints.py @@ -2,6 +2,7 @@ from copy import deepcopy from pathlib import Path +import torch from datasets import Dataset from bergson.collector.collector import ( @@ -24,6 +25,8 @@ setup_model_and_peft, ) +from .segment_aggregation import LAMBDA_COUNTS_FILENAME, write_lambda_counts + def precompute_checkpoint_averaged_lambdas( index_cfg: IndexConfig, @@ -49,7 +52,14 @@ def precompute_checkpoint_averaged_lambdas( eigen_path = base_run / f"segment_{seg}" / method out_path = ckpt_method_dir / output_subdir - if out_path.exists(): + # A run predating `fisher_normalization` has lambdas but no counts.json; + # recover the counts without redoing the data pass. + counts_only = ( + out_path.exists() + and resume + and not (out_path / LAMBDA_COUNTS_FILENAME).exists() + ) + if out_path.exists() and not counts_only: if resume: logger.info( f"[seg {seg} ckpt {idx_in_seg}] skip — exists at {out_path}" @@ -63,8 +73,13 @@ def precompute_checkpoint_averaged_lambdas( ) logger.info( - f"[seg {seg} ckpt {idx_in_seg}] computing {output_subdir} at " - f"model={ckpt!r} using eigvecs from {eigen_path}" + f"[seg {seg} ckpt {idx_in_seg}] " + + ( + f"backfilling {LAMBDA_COUNTS_FILENAME} for existing {output_subdir}" + if counts_only + else f"computing {output_subdir} at model={ckpt!r} " + f"using eigvecs from {eigen_path}" + ) ) ckpt_index_cfg = deepcopy(index_cfg) @@ -77,18 +92,31 @@ def precompute_checkpoint_averaged_lambdas( ds, _ = setup_data_pipeline(ckpt_index_cfg) - launch_distributed_run( - "checkpoint_averaged_lambda", - _lambda_worker, - [ - ckpt_index_cfg, - hessian_cfg, - ds, - ckpt_method_dir, - eigen_path, - output_subdir, - ], - ckpt_index_cfg.distributed, + if not counts_only: + launch_distributed_run( + "checkpoint_averaged_lambda", + _lambda_worker, + [ + ckpt_index_cfg, + hessian_cfg, + ds, + ckpt_method_dir, + eigen_path, + output_subdir, + ], + ckpt_index_cfg.distributed, + ) + + write_lambda_counts( + out_path, + documents=len(ds), + tokens=int( + torch.load( + ckpt_index_cfg.partial_run_path / "total_processed.pt", + map_location="cpu", + weights_only=False, + ) + ), ) @@ -151,10 +179,13 @@ def precompute_checkpoint_hessians( hessian_cfg: HessianConfig, approx_unrolling_cfg: ApproxUnrollingConfig, *, - overwrite: bool = False, + resume: bool = False, ) -> None: """Run :func:`approximate_hessians` once per checkpoint. + ``resume`` skips checkpoints whose output already exists, matching the + other steps of the pipeline. + The pipeline-level divisibility check (``n_ckpts % n_segments == 0``) runs in :func:`approx_unrolling_pipeline` before this is called. """ @@ -173,7 +204,7 @@ def precompute_checkpoint_hessians( out_path = ckpt_dir / method if out_path.exists(): - if not overwrite: + if resume: logger.info( f"[seg {seg} ckpt {idx_in_seg}] skip — exists at {out_path}" ) diff --git a/bergson/approx_unrolling/segment_aggregation.py b/bergson/approx_unrolling/segment_aggregation.py index fcfdbb76..eb0650cb 100644 --- a/bergson/approx_unrolling/segment_aggregation.py +++ b/bergson/approx_unrolling/segment_aggregation.py @@ -1,5 +1,7 @@ +import json import shutil from pathlib import Path +from typing import Literal import torch import torch.distributed as dist @@ -14,18 +16,62 @@ _SHARD_KINDS = ("activation_sharded", "gradient_sharded") +LAMBDA_COUNTS_FILENAME = "counts.json" + +FisherNormalization = Literal["document", "token", "none"] + + +def write_lambda_counts(lambda_dir: Path, *, documents: int, tokens: int) -> None: + """Record the denominators available for normalizing a checkpoint's lambdas. + + ``LambdaCollector`` accumulates a sum over documents; ``total_processed`` + counts tokens. Both are recorded so the choice is made at aggregation time. + """ + with open(lambda_dir / LAMBDA_COUNTS_FILENAME, "w") as f: + json.dump({"documents": documents, "tokens": tokens}, f) + + +def lambda_denominator( + input_dirs: list[Path], normalization: FisherNormalization +) -> float: + """Pooled count over a segment's checkpoints. + + The lambdas are summed over checkpoints, so the denominator sums too: the + result is the mean over every (checkpoint, document) pair, which is what + kronfluence's ``lambda_matrix / num_lambda_processed`` computes for a single + checkpoint. + """ + if normalization == "none": + return 1.0 + + key = "documents" if normalization == "document" else "tokens" + total = 0 + for d in input_dirs: + counts_path = d / LAMBDA_COUNTS_FILENAME + if not counts_path.exists(): + raise FileNotFoundError( + f"Missing {counts_path}, needed to normalize the segment " + f"eigenvalues by {key}. It is written alongside the lambda " + "shards; re-run the per-checkpoint lambda step, or set " + "`fisher_normalization: none` to keep the unnormalized sum." + ) + with open(counts_path) as f: + total += json.load(f)[key] + return float(total) + def sum_sharded_dirs( input_dirs: list[Path], output_dir: Path, distributed: DistributedConfig, + divisor: float = 1.0, ) -> None: """Sum per-rank shards across ``input_dirs`` into ``output_dir``.""" output_dir.mkdir(parents=True, exist_ok=True) launch_distributed_run( "sum_sharded_dirs", _sum_sharded_dirs_worker, - [input_dirs, output_dir], + [input_dirs, output_dir, divisor], distributed, ) @@ -36,12 +82,13 @@ def _sum_sharded_dirs_worker( world_size: int, input_dirs: list[Path], output_dir: Path, + divisor: float, ) -> None: init_dist(rank, local_rank, world_size) device = get_device(local_rank) shard_name = f"shard_{rank}.safetensors" in_paths = [d / shard_name for d in input_dirs] - _sum_my_shard(in_paths, output_dir / shard_name, device=device) + _sum_my_shard(in_paths, output_dir / shard_name, device=device, divisor=divisor) if world_size > 1: dist.barrier() @@ -50,6 +97,7 @@ def _sum_my_shard( in_paths: list[Path], out_path: Path, device: str, + divisor: float = 1.0, ) -> None: """Sum all tensor dicts in ``in_paths`` and write to ``out_path``.""" acc: dict[str, torch.Tensor] = {} @@ -61,6 +109,9 @@ def _sum_my_shard( acc[k] = t.clone() else: acc[k].add_(t) + if divisor != 1.0: + for v in acc.values(): + v.div_(divisor) save_file({k: v.cpu() for k, v in acc.items()}, out_path) @@ -191,10 +242,16 @@ def aggregate_segment_lambdas( distributed: DistributedConfig, *, resume: bool = False, + normalization: FisherNormalization = "document", input_subdir: str = "averaged_ev_correct_sharded", output_subdir: str = "eigenvalue_correction_sharded", ) -> None: - """Sum per-checkpoint lambdas into per-segment lambda.""" + """Sum per-checkpoint lambdas into per-segment lambda. + + ``normalization`` divides the sum by the pooled document or token count, so + the segment eigenvalues are an expected Fisher rather than a total that + grows with dataset size. See ``ApproxUnrollingConfig.fisher_normalization``. + """ logger = get_logger("aggregate_segment_lambdas") base_run = Path(run_path) @@ -217,5 +274,9 @@ def aggregate_segment_lambdas( f"Missing per-ckpt lambda dir {d}; did step 3 finish?" ) - logger.info(f"[seg {seg}] summing lambdas -> {out_dir}") - sum_sharded_dirs(input_dirs, out_dir, distributed) + divisor = lambda_denominator(input_dirs, normalization) + logger.info( + f"[seg {seg}] summing lambdas -> {out_dir} " + f"(normalization={normalization}, divisor={divisor:g})" + ) + sum_sharded_dirs(input_dirs, out_dir, distributed, divisor=divisor) diff --git a/bergson/approx_unrolling/trainer_run.py b/bergson/approx_unrolling/trainer_run.py index 5e37dd9d..b5a8198e 100644 --- a/bergson/approx_unrolling/trainer_run.py +++ b/bergson/approx_unrolling/trainer_run.py @@ -40,8 +40,10 @@ def write_lr_history( def load_training_config(trainer_run: str | Path) -> TrainingConfig: - """Load the ``TrainingConfig`` a run was launched with. ``save_run_config`` - writes ``{command_name: {...}}``, so the single value is the config.""" + """Load the ``TrainingConfig`` a run was launched with. + + ``save_run_config`` writes ``{steps: [{command_name: {...}}], metadata: ...}``, + so the training config is the first step's single value.""" path = Path(trainer_run) / CONFIG_FILENAME if not path.is_file(): raise FileNotFoundError( @@ -50,9 +52,11 @@ def load_training_config(trainer_run: str | Path) -> TrainingConfig: ) with open(path) as f: - loaded = yaml.safe_load(f) + loaded: Any = yaml.safe_load(f) - # One-step configs are a list of {command: payload}; take the first payload. + if isinstance(loaded, dict) and "steps" in loaded: + loaded = loaded["steps"] + # Steps are a list of {command: payload}; take the first payload. if isinstance(loaded, list): if not loaded: raise ValueError(f"{path} is empty") @@ -64,7 +68,12 @@ def load_training_config(trainer_run: str | Path) -> TrainingConfig: if not isinstance(payload, dict): raise ValueError(f"{path} is not a bergson run config") - return TrainingConfig.from_dict(payload, drop_extra_fields=True) + try: + return TrainingConfig.from_dict(payload, drop_extra_fields=True) + except Exception as e: + # An attribution run's config.yaml has the same shape, so reaching here + # is expected; callers that guess at a run dir catch ValueError. + raise ValueError(f"{path} is not a bergson training config: {e}") from e def derive_momentum(training_cfg: TrainingConfig) -> float: @@ -127,7 +136,17 @@ def resolve(cfg: ApproxUnrollingConfig) -> ApproxUnrollingConfig: cfg.momentum = 0.0 return cfg - training_cfg = load_training_config(trainer_run) + try: + training_cfg = load_training_config(trainer_run) + except ValueError as e: + # A directory with a config.yaml is not necessarily a training run: an + # attribution run writes one next to the checkpoints it produced. Infer + # nothing from it rather than failing the pipeline. + logger.warning("Ignoring %s as a trainer run: %s", trainer_run, e) + if cfg.momentum is None: + cfg.momentum = 0.0 + return cfg + filled: list[str] = [] if cfg.model_path is None: diff --git a/bergson/config/config.py b/bergson/config/config.py index 0730f480..5a4333fc 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -802,6 +802,26 @@ class ApproxUnrollingConfig(Serializable): Bae et al. 2024, Appendix C). Requires an ``optimizer.pt`` in every checkpoint dir; see :mod:`bergson.approx_unrolling.adam_preconditioner`.""" + adam_segment_hybrid: bool = False + """With ``use_adam_preconditioner``: use the Eq-43 hybrid for F_segment -- + the diagonal supplies only the matrix exponential, the EK-FAC factors supply + H^-1 (Bae et al. 2024, App. D: "we still use the EK-FAC factors to compute + H^-1 g in Equation 43").""" + + adam_hybrid_damping: float = 0.1 + """Relative damping for the hybrid's EK-FAC inverse.""" + + fisher_normalization: Literal["document", "token", "none"] = "document" + """Denominator for the segment EK-FAC eigenvalues, which set the scale of + ``lr*steps * sigma`` that ``f_segment``/``f_backward`` are evaluated at. + + - "document": divide by the number of documents summed over the segment's + checkpoints. Matches kronfluence (Bae et al.'s implementation), which + divides its lambda matrix by a per-sample count. + - "token": divide by the token count instead. + - "none": leave the raw sum. Pre-normalization behaviour; the resulting + scale grows with dataset size and checkpoints per segment.""" + query: DataConfig = field(default_factory=DataConfig) """Query dataset spec; gradients computed at the final checkpoint.""" diff --git a/bergson/hessians/apply_hessian.py b/bergson/hessians/apply_hessian.py index ed17fb28..f9a334fa 100644 --- a/bergson/hessians/apply_hessian.py +++ b/bergson/hessians/apply_hessian.py @@ -45,6 +45,14 @@ class EkfacConfig: diagonal approximation of P^1/2 H P^1/2 in parameter space (the Adam/AdamW approximate-unrolling variant) instead of on the eigenvalues in the EKFAC eigenbasis. Requires ``apply_fn``.""" + preconditioner_hybrid: bool = False + """With ``preconditioner_path``: instead of evaluating ``apply_fn`` wholly on + the diagonal, apply the elementwise ``apply_fn(p * diag(H))`` mask and then + the standard EK-FAC inverse. This is the Eq-43 reading of the Adam SOURCE + variant (Bae et al. 2024, App. C/D), which keeps eigenbasis curvature in the + inverse rather than approximating it diagonally.""" + preconditioner_hybrid_damping: float = 0.1 + """Relative damping for the hybrid's EK-FAC inverse.""" preconditioner_post_multiply: bool = False """With ``preconditioner_path``: multiply the applied function's output by the preconditioner grid — the P^1/2 F(M) P^1/2 sandwich of the segment @@ -97,10 +105,11 @@ def __init__( self.device = get_device(self.rank) def compute_ivhp_sharded(self): + chain: list = [] if self.cfg.preconditioner_path: if self.apply_fn is None: raise ValueError("preconditioner_path requires apply_fn.") - preconditioner = DiagonalFactoredPreconditioner.from_shards( + diagonal = DiagonalFactoredPreconditioner.from_shards( self.path, self.cfg.preconditioner_path, rank=self.rank, @@ -109,6 +118,25 @@ def compute_ivhp_sharded(self): multiply_by_preconditioner=self.cfg.preconditioner_post_multiply, ev_correction=self.cfg.ev_correction, ) + if self.cfg.preconditioner_hybrid: + # Bae et al. App. D: "use the diagonal Hessian approximation for + # computing the matrix exponential ... Note that we still use the + # EK-FAC factors to compute H^-1 g in Equation 43." + # Eq-43 reads M_mask @ H^-1 @ g_train; applied to the QUERY + # gradient that is the adjoint H^-1 @ M_mask @ q, so the + # diagonal mask goes first and the EK-FAC inverse second. + chain.append(diagonal) + preconditioner = FactoredPreconditioner.from_shards( + self.path, + rank=self.rank, + device=self.device, + inversion_cfg=InversionConfig( + damping_factor=self.cfg.preconditioner_hybrid_damping + ), + ev_correction=self.cfg.ev_correction, + ) + else: + preconditioner = diagonal else: preconditioner = FactoredPreconditioner.from_shards( self.path, @@ -173,6 +201,8 @@ def compute_ivhp_sharded(self): device=self.device, dtype=torch.float32 ) + for pre in chain: + grads = pre.apply(grads) transformed = preconditioner.apply(grads) del grads diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index e6a134da..d25eb87c 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -262,6 +262,7 @@ def worker( if getattr(run_cfg, "save_optimizer_state", "none") == "all" else None ), + save_interval=getattr(run_cfg, "save_interval", 0), ) # Called on every rank: FSDP moments are DTensors whose gather is a # collective; rank 0 writes inside. @@ -323,7 +324,12 @@ def worker( ) multi_query = False - if not score_path: + # Plain runs compute scores here; skip_metagradient turns MAGIC into a + # trainer, which is what the SOURCE/EK-FAC replications want from it. + do_metagradient = not score_path and not getattr( + run_cfg, "skip_metagradient", False + ) + if do_metagradient: # Sanity check if not isinstance(run_cfg, MagicConfig): raise RuntimeError("run_cfg must be a MagicConfig to compute scores") @@ -374,6 +380,19 @@ def worker( score_path = os.path.join(run_cfg.run_path, "scores.pt") torch.save(scores, score_path) print(f"Saved attribution scores to {score_path}") + elif not score_path: + # skip_metagradient: no MAGIC scores; build the bank with dummy zeros. + scores = torch.zeros_like(stream.weights).cpu() + if pad_count: + scores = ( + scores[:-weight_pad_count] if scores.ndim == 1 else scores[:-pad_count] + ) + if global_rank == 0: + print(f"Baseline loss: {baseline}") + print( + "skip_metagradient: skipped MAGIC backward; using dummy zero " + "scores (bank build only)." + ) elif os.path.isdir(score_path) or score_path.endswith(".npy"): scores, multi_query = load_attribution_scores(score_path) else: diff --git a/bergson/magic/config.py b/bergson/magic/config.py index 2fa0f944..b95a0bdd 100644 --- a/bergson/magic/config.py +++ b/bergson/magic/config.py @@ -3,7 +3,7 @@ from ..config.config import ValidationConfig -MagicSaveMode = Literal["all", "sqrt", "log"] +MagicSaveMode = Literal["all", "sqrt", "log", "interval"] @dataclass @@ -20,10 +20,17 @@ class MagicConfig(ValidationConfig): - 'sqrt' saves at a linearly-spaced interval, every sqrt(N) steps. This method uses O(sqrt N) space and O(N) time. + - 'interval' saves every `save_interval` steps (plus the final state when the + cadence lands on it). Use for SOURCE-style runs that need a few evenly + spaced checkpoints rather than backward-replay coverage. + The original MAGIC paper used 'log', but 'sqrt' is often a better choice when disk space is not a concern. """ + save_interval: int = 0 + """Snapshot spacing in steps for `save_mode: interval`.""" + backward_save_every: int = 0 """How often (in steps) to save backward state for resume.""" @@ -36,3 +43,9 @@ class MagicConfig(ValidationConfig): skip_validation: bool = False """Stop after computing and saving attribution scores, before the leave-k-out retraining loop. Useful for score-only MAGIC runs.""" + + skip_metagradient: bool = False + """Skip the MAGIC backward pass entirely, using the trainer only. Attribution + scores are filled with zeros. Useful when another method (SOURCE, EK-FAC) + supplies the scores and MAGIC is wanted purely for its checkpoint schedule + and optimizer-state export.""" diff --git a/bergson/magic/metasmoothness.py b/bergson/magic/metasmoothness.py index e96adc78..1e20beba 100644 --- a/bergson/magic/metasmoothness.py +++ b/bergson/magic/metasmoothness.py @@ -46,6 +46,39 @@ def metasmoothness_score( return float((d / total * s1 * s2).sum()) +def metasmoothness_group_scores( + theta0: torch.Tensor, + theta_h: torch.Tensor, + theta_2h: torch.Tensor, + group_ids: torch.Tensor, + group_names: list[str], +) -> dict: + """Per-optimizer-path breakdown of the metasmoothness score. + + Muon routes 2D parameters through Newton-Schulz and 1D parameters through + an AdamW fallback; ``eps_root`` only enters the AdamW branch. Since the + global score weights coordinates by their L1 movement, the 2D group (99.9% + of GPT-2's parameters) can dominate it and mask a non-smooth 1D group. + + The decomposition is exact: ``score == sum(share_g * score_g)``. + """ + d = (theta_2h - theta0).abs() + total = d.sum() + contrib = d * torch.sign(theta_h - theta0) * torch.sign(theta_2h - theta_h) + + out = {} + for gid, name in enumerate(group_names): + m = group_ids == gid + dg = d[m].sum() + out[name] = { + "score": float(contrib[m].sum() / dg) if dg > 0 else 1.0, + "movement_share": float(dg / total) if total > 0 else 0.0, + "movement_l1": float(dg), + "numel": int(m.sum()), + } + return out + + def metasmoothness_worker( global_rank: int, rank: int, @@ -97,6 +130,9 @@ def metasmoothness_worker( v[-weight_pad_count:] = 0.0 thetas: list[torch.Tensor] = [] + theta_init: torch.Tensor | None = None + group_ids: torch.Tensor | None = None + group_names: list[str] = [] for k in range(3): weights = 1.0 + run_cfg.fd_step * k * v if pad_count: @@ -108,6 +144,16 @@ def metasmoothness_worker( trainer, fwd_state, model = prepare_trainer(run_cfg, rank, schedule) fwd_state.detach_() + if k == 0 and global_rank == 0: + theta_init = torch.cat( + [p.detach().float().cpu().flatten() for p in fwd_state.params.values()] + ) + if run_cfg.save_models: + os.makedirs(run_cfg.run_path, exist_ok=True) + torch.save( + {n: p.detach().float().cpu() for n, p in fwd_state.params.items()}, + os.path.join(run_cfg.run_path, "theta_init.pt"), + ) fwd_state = trainer.train(fwd_state, stream, inplace=True, fsdp=run_cfg.fsdp) if global_rank == 0: @@ -115,6 +161,27 @@ def metasmoothness_worker( [p.detach().float().cpu().flatten() for p in fwd_state.params.values()] ) thetas.append(theta) + if group_ids is None and run_cfg.optimizer == "muon": + # Muon's own split: ndim==2 -> Newton-Schulz, else AdamW fallback + # (the only branch eps_root enters). See optim.py muon(). + group_names = ["muon_2d", "adamw_1d"] + group_ids = torch.cat( + [ + torch.full( + (p.numel(),), 0 if p.ndim == 2 else 1, dtype=torch.uint8 + ) + for p in fwd_state.params.values() + ] + ) + if run_cfg.save_models and k == 0: + # Only the unperturbed run: the k>0 models are finite-difference + # probes, not models anyone wants to reuse. + path = os.path.join(run_cfg.run_path, "theta_final.pt") + torch.save( + {n: p.detach().float().cpu() for n, p in fwd_state.params.items()}, + path, + ) + print(f"[metasmoothness] saved {path}") print(f"[metasmoothness] finished training {k + 1}/3 (w = 1 + {k}*h*v)") del trainer, fwd_state, model @@ -127,7 +194,27 @@ def metasmoothness_worker( "direction_seed": run_cfg.direction_seed, "total_movement_l1": movement, } + if theta_init is not None: + # Relative update norms of the unperturbed run, matching the ΔL1/ΔL2 + # convention in LDS_RESULTS (‖θ_final − θ_init‖ / ‖θ_init‖). + delta = thetas[0] - theta_init + result["delta_l1"] = float(delta.abs().sum() / theta_init.abs().sum()) + result["delta_l2"] = float(delta.norm() / theta_init.norm()) + print( + f"[metasmoothness] deltaL1 = {result['delta_l1']:.4g} " + f"deltaL2 = {result['delta_l2']:.4g}" + ) print(f"[metasmoothness] score = {score:.4f} (h={run_cfg.fd_step})") + if group_ids is not None: + groups = metasmoothness_group_scores( + thetas[0], thetas[1], thetas[2], group_ids, group_names + ) + result["groups"] = groups + for name, g in groups.items(): + print( + f"[metasmoothness] {name}: score = {g['score']:.4f} " + f"share = {g['movement_share']:.4f} numel = {g['numel']:,}" + ) os.makedirs(run_cfg.run_path, exist_ok=True) with open(os.path.join(run_cfg.run_path, "metasmoothness.json"), "w") as f: json.dump(result, f, indent=2) diff --git a/bergson/magic/trainer.py b/bergson/magic/trainer.py index 6a92c33a..9acf7015 100644 --- a/bergson/magic/trainer.py +++ b/bergson/magic/trainer.py @@ -106,7 +106,9 @@ def result(self): return result -def next_save_index(current: int, n: int, save_mode: MagicSaveMode) -> int: +def next_save_index( + current: int, n: int, save_mode: MagicSaveMode, save_interval: int = 0 +) -> int: """Index of the checkpoint following the one saved at step `current`. `n` is the total number of training steps. The result is always strictly @@ -117,6 +119,10 @@ def next_save_index(current: int, n: int, save_mode: MagicSaveMode) -> int: case "all": # Save every step return current + 1 + case "interval": + if save_interval <= 0: + raise ValueError("save_mode='interval' requires save_interval > 0.") + return current + save_interval case "sqrt": chunk_size = math.isqrt(n) @@ -555,6 +561,7 @@ def train( max_grad_norm: float | None = None, grad_accum_steps: int = 1, optimizer_cfg: dict | None = None, + save_interval: int = 0, ) -> TrainerState: """Train the model on the given data stream, starting from the given state. @@ -583,6 +590,7 @@ def train( optimizer_cfg: When set (to the optimizer's betas/eps/eps_root), write each checkpoint's second moments to ``optimizer.pt``, tagged with the step and these hyperparameters. AdamW only. + save_interval: Snapshot spacing in steps for ``save_mode="interval"``. Returns: The final trainer state after training. @@ -634,7 +642,7 @@ def train( pending_save = state.save(p, debug_pbar=pbar if debug else None) - next_save = next_save_index(next_save, n, save_mode) + next_save = next_save_index(next_save, n, save_mode, save_interval) x = data[i] state = self.step( @@ -653,6 +661,23 @@ def train( if pending_save is not None: pending_save.result() + # Snapshots are written BEFORE each step, so the post-training state has + # no snapshot of its own; write it when the cadence lands exactly on n. + # Interval mode only: the other modes' backward replay indexes the data + # stream by checkpoint, so an extra trailing snapshot runs it past the end. + if save_dir and save_mode == "interval" and next_save == n: + p = os.path.join(save_dir, f"step_{n}.ckpt") + if optimizer_cfg is not None: + os.makedirs(p, exist_ok=True) + save_second_moments_as_optimizer_pt( + self.model, + state.opt_state, + os.path.join(p, "optimizer.pt"), + step=n, + **optimizer_cfg, + ) + state.save(p).result() + return state def save_backward_state(self, bwd_state, path, expected_idx, last_idx): diff --git a/bergson/validate.py b/bergson/validate.py index 447613d2..ba824707 100644 --- a/bergson/validate.py +++ b/bergson/validate.py @@ -80,6 +80,62 @@ def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]: return scores, False +def lds_from_precomputed_subsets( + score_path: str, + masks_path: str, + losses_path: str, + summary_path: str | None = None, +) -> float: + """LDS against a precomputed subset-retrain ground truth. + + ``masks_path``: ``[M, N_train]`` 0/1 subset-membership array (row ``j`` + marks the training rows the ``j``-th bank model was trained ON); + ``losses_path``: ``[M, Q]`` per-subset query losses (averaged over retrain + seeds). Follows Park et al. 2023 / Bae et al. 2024 Eq. 26-27: the group + prediction for subset ``j`` is the summed score of its members, and the LDS + is the per-query Spearman between predictions and subset losses, averaged + over queries. ``load_attribution_scores`` orients scores to the loss-diff + convention, under which the prediction correlates POSITIVELY with the + subset loss, so the returned mean Spearman is comparable to published LDS. + """ + scores, _ = load_attribution_scores(score_path) + if scores.ndim == 1: + scores = scores[:, None] + if scores.ndim != 2: + raise ValueError( + f"Expected per-document scores [N_train, Q]; got shape " + f"{tuple(scores.shape)}." + ) + masks = np.asarray(torch.load(masks_path, weights_only=False), dtype=np.float64) + losses = np.asarray(torch.load(losses_path, weights_only=False), dtype=np.float64) + if masks.shape[1] != scores.shape[0]: + raise ValueError( + f"masks cover {masks.shape[1]} training rows but scores have " + f"{scores.shape[0]}; the score run must use the identical dataset " + "ordering the masks index into." + ) + if losses.shape != (masks.shape[0], scores.shape[1]): + raise ValueError( + f"losses shape {losses.shape} does not match " + f"[num_subsets={masks.shape[0]}, num_queries={scores.shape[1]}]." + ) + + preds = masks @ scores.double().numpy() + per_query = np.array( + [spearmanr(preds[:, q], losses[:, q]).statistic for q in range(losses.shape[1])] + ) + mean_lds = float(np.nanmean(per_query)) + + if summary_path is not None: + writer = CSVWriter(summary_path, ["query", "spearman_corr"]) + for q, r in enumerate(per_query): + writer.writerow(q, float(r)) + writer.writerow("mean", mean_lds) + writer.close() + print(f"LDS (mean Spearman over {losses.shape[1]} queries): {mean_lds:.4f}") + return mean_lds + + def bank_loss_cache_key( run_cfg: ValidationConfig, multi_query: bool, num_subsets: int ) -> str: diff --git a/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm.yaml b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm.yaml new file mode 100644 index 00000000..d6303c17 --- /dev/null +++ b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm.yaml @@ -0,0 +1,83 @@ +# SOURCE (approximate unrolling, Bae et al. 2024) for the lotus GPT-2 base +# model against the 50 test queries, validated on the 400-model lotus bank. +# Plain SGD variant (no optimizer preconditioning). +# +# fisher_normalization: document -- the segment EK-FAC eigenvalues are a mean +# over the segment's (checkpoint, document) pairs, matching kronfluence's +# lambda_matrix / num_lambda_processed. This is the replication setting. +# +# Baseline for comparison: runs/lotus_source_q50_damp0 (unnormalized sums, +# LDS rho=0.387 +/- 0.039). + +# Run with `bergson examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm.yaml` + +steps: + - approxunrolling: + index_cfg: + run_path: runs/lotus_source_q50_docnorm + # Tokenizer + final-model reference; per-checkpoint passes load the + # model dirs below. + model: runs/lotus/retrained/base + # In the approxunrolling pipeline `overwrite` doubles as resume: + # steps whose output dirs exist are skipped after a crash. + overwrite: true + precision: fp32 + # GPT-2's max context length is 1024. + token_batch_size: 1024 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 # <-- adjust to this node's GPU count + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + # Doc-level scoring (one doc per row, no packing): packed 512-token + # chunks contaminate per-doc gradients across doc boundaries and + # were found to suppress LDS (see gpt2_lotus_trackstar50q_docs). + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + approx_unrolling_cfg: + checkpoints: + - runs/lotus_source_q50/models/step_51 + - runs/lotus_source_q50/models/step_102 + - runs/lotus_source_q50/models/step_136 + - runs/lotus_source_q50/models/step_187 + - runs/lotus_source_q50/models/step_238 + - runs/lotus_source_q50/models/step_288 + segments: 3 + fisher_normalization: document + lr_list: [4.857423e-04, 5.616667e-04, 2.416667e-04] + step_size_list: [97, 96, 96] + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + # One score column per query (50 total), like the EK-FAC and + # TrackStar runs, for multi-query LDS validation. + query_aggregation: none + + - validate: + run_path: runs/lotus_source_q50_docnorm_validate + model: gpt2 + overwrite: true + + scores: runs/lotus_source_q50_docnorm/scores + retrained_dir: runs/lotus + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 512 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + + batch_size: 64 diff --git a/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm_validate.yaml b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm_validate.yaml new file mode 100644 index 00000000..bda9d939 --- /dev/null +++ b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm_validate.yaml @@ -0,0 +1,18 @@ +# Validation only: the SOURCE scores for the per-document run already exist. + +steps: +- validate: + run_path: runs/lotus_source_q50_docnorm_validate + model: gpt2 + overwrite: true + scores: runs/lotus_source_q50_docnorm/scores + retrained_dir: runs/lotus + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: train + chunk_length: 512 + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: test[1:51] + chunk_length: 0 + batch_size: 64 diff --git a/examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml new file mode 100644 index 00000000..9568a898 --- /dev/null +++ b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml @@ -0,0 +1,83 @@ +# SOURCE (approximate unrolling, Bae et al. 2024) for the lotus GPT-2 base +# model against the 50 test queries, validated on the 400-model lotus bank. +# Plain SGD variant (no optimizer preconditioning). +# +# fisher_normalization: token -- as the document variant, but dividing by the +# token count instead. LambdaCollector sums over documents, so this is the +# per-document mean scaled by ~1/mean_tokens_per_doc. +# +# Baseline for comparison: runs/lotus_source_q50_damp0 (unnormalized sums, +# LDS rho=0.387 +/- 0.039). + +# Run with `bergson examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml` + +steps: + - approxunrolling: + index_cfg: + run_path: runs/lotus_source_q50_toknorm + # Tokenizer + final-model reference; per-checkpoint passes load the + # model dirs below. + model: runs/lotus/retrained/base + # In the approxunrolling pipeline `overwrite` doubles as resume: + # steps whose output dirs exist are skipped after a crash. + overwrite: true + precision: fp32 + # GPT-2's max context length is 1024. + token_batch_size: 1024 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 # <-- adjust to this node's GPU count + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + # Doc-level scoring (one doc per row, no packing): packed 512-token + # chunks contaminate per-doc gradients across doc boundaries and + # were found to suppress LDS (see gpt2_lotus_trackstar50q_docs). + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + approx_unrolling_cfg: + checkpoints: + - runs/lotus_source_q50/models/step_51 + - runs/lotus_source_q50/models/step_102 + - runs/lotus_source_q50/models/step_136 + - runs/lotus_source_q50/models/step_187 + - runs/lotus_source_q50/models/step_238 + - runs/lotus_source_q50/models/step_288 + segments: 3 + fisher_normalization: token + lr_list: [4.857423e-04, 5.616667e-04, 2.416667e-04] + step_size_list: [97, 96, 96] + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + # One score column per query (50 total), like the EK-FAC and + # TrackStar runs, for multi-query LDS validation. + query_aggregation: none + + - validate: + run_path: runs/lotus_source_q50_toknorm_validate + model: gpt2 + overwrite: true + + scores: runs/lotus_source_q50_toknorm/scores + retrained_dir: runs/lotus + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 512 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + + batch_size: 64 diff --git a/examples/replication/wikitext_gpt2_ekfac.yaml b/examples/replication/wikitext_gpt2_ekfac.yaml new file mode 100644 index 00000000..59de6aa1 --- /dev/null +++ b/examples/replication/wikitext_gpt2_ekfac.yaml @@ -0,0 +1,56 @@ +# EK-FAC influence functions at the final checkpoint of the WikiText-2 +# replication run (examples/replication/wikitext_gpt2_train.yaml), scored +# against all 481 validation queries. +# +# Reference: kronfluence examples/wikitext reports LDS 0.44 against its shipped +# ground truth (masks.pt 100x4656 alpha=0.5 keep-masks, losses.pt 100x481 +# five-seed-averaged subset losses). Running this pipeline on kronfluence's own +# released model gives 0.4431, so the EK-FAC implementation itself is validated; +# this config checks that a model trained by OUR trainer lands in the same place. +# +# damping_factor 1e-8 is effectively undamped, matching the reference +# implementation, which adds damping only via its heuristic when unset. +# +# Run with: +# PYTHONPATH=$PWD python -m bergson examples/replication/wikitext_gpt2_ekfac.yaml + +steps: + - ekfac: + index_cfg: + run_path: runs/wikitext_repro/if_scores + model: runs/wikitext_repro/train/retrained/base + overwrite: true + precision: fp32 + token_batch_size: 1024 + max_batch_size: 64 + # GPT-2 ties lm_head to wte; tracking both double-counts the embedding. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 + nnode: 1 + data: + dataset: runs/wikitext_repro/data + split: train + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + score_cfg: + batch_size: 1024 + query_batch_size: 64 + + preprocess_cfg: + unit_normalize: false + + hessian_pipeline_cfg: + query: + dataset: runs/wikitext_repro/data + split: validation + chunk_length: 0 + # One score column per validation query, for per-query LDS. + query_aggregation: none + inversion_cfg: + damping_factor: 1.0e-8 diff --git a/examples/replication/wikitext_gpt2_source.yaml b/examples/replication/wikitext_gpt2_source.yaml new file mode 100644 index 00000000..b118dd47 --- /dev/null +++ b/examples/replication/wikitext_gpt2_source.yaml @@ -0,0 +1,76 @@ +# SOURCE on the WikiText-2 replication run, following Bae et al. 2024. +# +# The paper's WikiText-2/GPT-2 experiment is AdamW only (App. B.1, p. 28), so +# this uses the preconditioned variant of App. C, with the Eq-43 hybrid from +# App. D: "use the diagonal Hessian approximation for computing the matrix +# exponential in Equation 39 and Equation 43. Note that we still use the EK-FAC +# factors to compute H^-1 g in Equation 43." +# +# L=3 segments over C=6 checkpoints (Figure 6 uses L=3), constant lr 3e-5, +# 582 steps per segment. Queries: all 481 validation chunks, one score column +# each, for per-query LDS against the same ground truth as the EK-FAC config. +# +# adam_hybrid_damping 1e-8 keeps the chained EK-FAC inverse effectively +# undamped, matching the reference implementation. +# +# fisher_normalization: document is required, not cosmetic. The eigenfunctions +# are not scale-invariant, and Eq. 1 defines the empirical risk as a mean over +# data points, so the segment eigenvalues must be a per-document expectation. +# Token units put the mask argument in its dead-linear range (LDS 0.125 vs +# 0.383 in an earlier run of this experiment). +# +# The checkpoint dirs come from exporting the trainer's DCP snapshots: +# python -c "from bergson.utils.trainer_export import export_checkpoints; \ +# export_checkpoints('runs/wikitext_repro/train', \ +# steps=[291,582,873,1164,1455,1746], overwrite=True)" +# +# Run with: +# PYTHONPATH=$PWD python -m bergson examples/replication/wikitext_gpt2_source.yaml + +steps: + - approxunrolling: + index_cfg: + run_path: runs/wikitext_repro/source_scores + model: runs/wikitext_repro/train/retrained/base + overwrite: true + precision: fp32 + token_batch_size: 1024 + max_batch_size: 64 + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 + nnode: 1 + data: + dataset: runs/wikitext_repro/data + split: train + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + approx_unrolling_cfg: + checkpoints: + - runs/wikitext_repro/train/exported/checkpoint-291 + - runs/wikitext_repro/train/exported/checkpoint-582 + - runs/wikitext_repro/train/exported/checkpoint-873 + - runs/wikitext_repro/train/exported/checkpoint-1164 + - runs/wikitext_repro/train/exported/checkpoint-1455 + - runs/wikitext_repro/train/exported/checkpoint-1746 + segments: 3 + lr_list: [3.0e-5, 3.0e-5, 3.0e-5] + step_size_list: [582, 582, 582] + query: + dataset: runs/wikitext_repro/data + split: validation + chunk_length: 0 + query_aggregation: none + query_batch_size: 64 + + # App. C preconditioned variant + the App. D Eq-43 hybrid. + use_adam_preconditioner: true + adam_segment_hybrid: true + adam_hybrid_damping: 1.0e-8 + + fisher_normalization: document diff --git a/examples/replication/wikitext_gpt2_train.yaml b/examples/replication/wikitext_gpt2_train.yaml new file mode 100644 index 00000000..9d6e9130 --- /dev/null +++ b/examples/replication/wikitext_gpt2_train.yaml @@ -0,0 +1,67 @@ +# Bae et al. 2024 WikiText-2 / GPT-2 training recipe (Appendix B.1, p. 28): +# "we fine-tuned the GPT-2 model using the WikiText-2 dataset. We followed the +# training script from the Transformer library but set the maximum sequence +# length to 512. During fine-tuning with AdamW, we saved 6 intermediate +# checkpoints for data attribution. The learning rate, weight decay, and +# batch size were set to 3e-5, 1e-2, and 8, respectively." +# +# Matches kronfluence examples/wikitext (3 epochs, seed 1004) so the resulting +# model is comparable to the one their reported LDS numbers were computed on. +# +# train_mode: true is load-bearing. kronfluence's train.py calls model.train(), +# so GPT-2's default 0.1 resid/embd/attn dropout is active during fine-tuning. +# Training with it off produces a measurably different model: EK-FAC LDS 0.372 +# on a dropout-free model vs 0.443 on kronfluence's own released model. +# +# 4656 train rows / 8 per batch = 582 steps per epoch, 1746 total; a checkpoint +# every 291 steps gives the paper's 6. +# +# Run with: +# PYTHONPATH=$PWD python -m bergson examples/replication/wikitext_gpt2_train.yaml + +steps: + - magic: + run_path: runs/wikitext_repro/train + model: gpt2 + overwrite: true + precision: fp32 + seed: 1004 + + # MAGIC is used purely as a trainer here; no metagradient is needed. + skip_metagradient: true + skip_validation: true + cleanup_ckpts: false + + data: + dataset: runs/wikitext_repro/data + split: train + chunk_length: 0 + query: + dataset: runs/wikitext_repro/data + split: validation + chunk_length: 0 + + distributed: + nproc_per_node: 1 + nnode: 1 + + optimizer: adamw + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-8 + eps_root: 0.0 + weight_decay: 0.01 + batch_size: 8 + num_epochs: 3 + lr_schedule: + lr_scheduler_type: constant + lr: 3.0e-5 + + # Dropout active, as in kronfluence's train.py. + train_mode: true + + save_mode: interval + save_interval: 291 + # SOURCE's Adam variant needs the optimizer state at every checkpoint. + save_optimizer_state: all + save_models: true diff --git a/experiments/source_fisher_normalization/chain_toknorm.sh b/experiments/source_fisher_normalization/chain_toknorm.sh new file mode 100755 index 00000000..d4875d7c --- /dev/null +++ b/experiments/source_fisher_normalization/chain_toknorm.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# Wait for the docnorm SOURCE run to finish, then start the toknorm one. +# Both need all 8 GPUs, so they cannot overlap. +# +# Bounded: gives up at DEADLINE_S, and exits early if the docnorm process +# disappears without producing scores. +set -u + +REPO=/mnt/ssd-1/lucia/bergson-damping +cd "$REPO" || exit 1 + +DEADLINE_S=$((6 * 3600)) +DOC_SCORES=runs/lotus_source_q50_docnorm/scores.npy +DOC_VALIDATE=runs/lotus_source_q50_docnorm_validate/summary.csv +TOK_CFG=examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml + +start=$(date +%s) +while true; do + now=$(date +%s) + elapsed=$((now - start)) + if [ "$elapsed" -gt "$DEADLINE_S" ]; then + echo "CHAIN: deadline of ${DEADLINE_S}s reached; not launching toknorm" + exit 1 + fi + + if [ -f "$DOC_VALIDATE" ]; then + echo "CHAIN: docnorm validation complete after ${elapsed}s" + break + fi + + if ! pgrep -f "bergson.*docnorm" >/dev/null 2>&1; then + # Process gone. Only proceed if it actually got to the end. + if [ -f "$DOC_VALIDATE" ]; then + echo "CHAIN: docnorm finished after ${elapsed}s" + break + fi + echo "CHAIN: docnorm process gone after ${elapsed}s without $DOC_VALIDATE" \ + "(scores.npy exists: $([ -f "$DOC_SCORES" ] && echo yes || echo no)); aborting" + exit 1 + fi + + sleep 60 +done + +CMD="PYTHONPATH=$REPO python -m bergson $TOK_CFG" +echo "CHAIN: launching toknorm" +echo "CMD: $CMD" +echo "CMD: $CMD" > logs/source_toknorm.log +PYTHONPATH="$REPO" python -m bergson "$TOK_CFG" >> logs/source_toknorm.log 2>&1 +echo "CHAIN: toknorm exited with $?" diff --git a/experiments/source_fisher_normalization/compare_lds.py b/experiments/source_fisher_normalization/compare_lds.py new file mode 100644 index 00000000..17a812be --- /dev/null +++ b/experiments/source_fisher_normalization/compare_lds.py @@ -0,0 +1,81 @@ +"""Compare LDS across SOURCE Fisher-normalization settings. + +Reads each validate run's summary.csv (one row per query) and reports the mean +correlation with a 95% CI over queries, which is the convention behind the +published `0.387 +/- 0.039`. + + python experiments/source_fisher_normalization/compare_lds.py \ + none=runs/lotus_source_q50_damp0_validate \ + document=runs/lotus_source_q50_docnorm_validate \ + token=runs/lotus_source_q50_toknorm_validate +""" + +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy.stats import wilcoxon + +CI95 = 1.959963985 + + +def summarize(path: Path) -> dict: + df = pd.read_csv(path / "summary.csv") + out: dict = {"n_queries": len(df)} + for col, key in (("spearman_corr", "rho"), ("pearson_corr", "r")): + v = df[col].to_numpy() + out[key] = v.mean() + out[f"{key}_ci"] = CI95 * v.std(ddof=1) / np.sqrt(len(v)) + out[f"{key}_raw"] = v + return out + + +def main() -> None: + args = sys.argv[1:] + if not args: + raise SystemExit(__doc__) + + runs = {} + for arg in args: + label, _, path = arg.partition("=") + p = Path(path) / "summary.csv" + if not p.exists(): + print(f"skip {label}: no {p}") + continue + runs[label] = summarize(Path(path)) + + print(f"\n{'normalization':<14}{'rho':>18}{'r':>18}{'queries':>9}") + print("-" * 59) + for label, s in runs.items(): + print( + f"{label:<14}" + f"{s['rho']:>10.3f} +/- {s['rho_ci']:.3f}" + f"{s['r']:>10.3f} +/- {s['r_ci']:.3f}" + f"{s['n_queries']:>9}" + ) + + # Paired over the same 50 queries, so a paired test is the informative one. + if "none" in runs: + base = runs["none"] + for label, s in runs.items(): + if label == "none" or len(s["rho_raw"]) != len(base["rho_raw"]): + continue + delta = s["rho_raw"] - base["rho_raw"] + stat = wilcoxon(delta) + print( + f"\n{label} vs none: mean drho={delta.mean():+.3f} " + f"(won {int((delta > 0).sum())}/{len(delta)} queries, " + f"wilcoxon p={stat.pvalue:.2g})" + ) + + print("\nLaTeX rows:") + for label, s in runs.items(): + print( + f"SOURCE ({label}) & ${s['rho']:.3f} \\pm {s['rho_ci']:.3f}$ " + f"& ${s['r']:.3f} \\pm {s['r_ci']:.3f}$ \\\\" + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/source_fisher_normalization/setup_run.py b/experiments/source_fisher_normalization/setup_run.py new file mode 100644 index 00000000..62efa7f5 --- /dev/null +++ b/experiments/source_fisher_normalization/setup_run.py @@ -0,0 +1,107 @@ +"""Seed a SOURCE run dir from an existing one, reusing everything the Fisher +normalization does not change. + +Only the segment lambdas (step 4) and everything downstream depend on +``fisher_normalization``; the per-checkpoint covariances/lambdas, the segment +eigenvectors and the query gradients do not. Those are hardlinked from the +baseline run so a normalization sweep costs no extra Hessian passes. + +Derived outputs are deliberately NOT linked: the pipeline truncates them in +place, which would corrupt the baseline through a shared inode. + + python experiments/source_fisher_normalization/setup_run.py \ + runs/lotus_source_q50_damp0 runs/lotus_source_q50_docnorm +""" + +import argparse +import os +from pathlib import Path + +# Per-segment inputs that are independent of the lambda normalization. +SEGMENT_REUSE = ( + "activation_sharded", + "gradient_sharded", + "eigen_activation_sharded", + "eigen_gradient_sharded", + "total_processed.pt", +) +CKPT_REUSE = ( + "activation_sharded", + "gradient_sharded", + "averaged_ev_correct_sharded", + "total_processed.pt", +) + + +def link_tree(src: Path, dst: Path) -> int: + """Hardlink every file under ``src`` into ``dst``. Returns the file count.""" + n = 0 + if src.is_file(): + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + os.link(src, dst) + return 1 + for path in src.rglob("*"): + if path.is_dir(): + continue + out = dst / path.relative_to(src) + out.parent.mkdir(parents=True, exist_ok=True) + if not out.exists(): + os.link(path, out) + n += 1 + return n + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("baseline", type=Path) + ap.add_argument("target", type=Path) + args = ap.parse_args() + + baseline, target = args.baseline, args.target + if not baseline.exists(): + raise SystemExit(f"baseline run {baseline} not found") + if target.exists(): + raise SystemExit(f"{target} already exists; remove it first") + + total = 0 + total += link_tree(baseline / "query", target / "query") + + for seg_dir in sorted(baseline.glob("segment_*")): + seg = seg_dir.name + # A Hessian method dir, as opposed to a derived dir (scores, + # query_grad_*) that also happens to hold a total_processed.pt. + method_dirs = [ + d for d in seg_dir.iterdir() if (d / "eigen_activation_sharded").is_dir() + ] + for method_dir in method_dirs: + for name in SEGMENT_REUSE: + src = method_dir / name + if src.exists(): + total += link_tree(src, target / seg / method_dir.name / name) + + for ckpt_dir in sorted(seg_dir.glob("ckpt_*")): + for method_dir in ckpt_dir.iterdir(): + if not method_dir.is_dir(): + continue + # kfac.part holds the token count the backfill reads. + names = ( + CKPT_REUSE + if not method_dir.name.endswith(".part") + else ("total_processed.pt",) + ) + for name in names: + src = method_dir / name + if src.exists(): + total += link_tree( + src, target / seg / ckpt_dir.name / method_dir.name / name + ) + + print(f"hardlinked {total} files from {baseline} -> {target}") + print( + "not linked (regenerated): eigenvalue_correction_sharded, query_grad_*, scores" + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/wikitext_replication/run_chain.py b/experiments/wikitext_replication/run_chain.py new file mode 100644 index 00000000..ec87a8ea --- /dev/null +++ b/experiments/wikitext_replication/run_chain.py @@ -0,0 +1,64 @@ +"""Export checkpoints -> EK-FAC IF -> SOURCE -> LDS for the WikiText-2 replication. + +Assumes examples/replication/wikitext_gpt2_train.yaml has already run. +Ground truth is kronfluence's shipped masks/losses (see GROUND_TRUTH below). + + PYTHONPATH=$PWD python experiments/wikitext_replication/run_chain.py + +Last run (2026-07-31, 8x GPU): EK-FAC IF 0.4451 (kronfluence reports 0.44), +SOURCE 0.4506. SOURCE > IF reproduces the ordering in the paper's Figure 6. +""" + +import subprocess +import sys +from pathlib import Path + +from bergson.utils.trainer_export import export_checkpoints +from bergson.validate import lds_from_precomputed_subsets + +RUN = Path("runs/wikitext_repro") +GROUND_TRUTH = Path("/mnt/ssd-1/lucia/bergson-source-repro/runs/wikitext_kron") +STEPS = [291, 582, 873, 1164, 1455, 1746] + + +def run(cmd: str) -> None: + print(f"\n=== RUN: {cmd}", flush=True) + subprocess.run(cmd, shell=True, check=True) + + +def main() -> None: + masks, losses = GROUND_TRUTH / "masks.pt", GROUND_TRUTH / "losses.pt" + for p in (masks, losses): + if not p.exists(): + sys.exit(f"missing ground truth: {p}") + + exported = export_checkpoints(RUN / "train", steps=STEPS, overwrite=True) + print(f"exported: {[str(p) for p in exported]}", flush=True) + + run( + "python -m bergson examples/replication/wikitext_gpt2_ekfac.yaml" + f" > {RUN}/if.log 2>&1" + ) + print("EK-FAC IF: ", end="") + lds_from_precomputed_subsets( + str(RUN / "if_scores/scores"), + str(masks), + str(losses), + summary_path=str(RUN / "if_lds.csv"), + ) + + run( + "python -m bergson examples/replication/wikitext_gpt2_source.yaml" + f" > {RUN}/source.log 2>&1" + ) + print("SOURCE: ", end="") + lds_from_precomputed_subsets( + str(RUN / "source_scores/scores"), + str(masks), + str(losses), + summary_path=str(RUN / "source_lds.csv"), + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_lds_precomputed_subsets.py b/tests/test_lds_precomputed_subsets.py new file mode 100644 index 00000000..e24a9de1 --- /dev/null +++ b/tests/test_lds_precomputed_subsets.py @@ -0,0 +1,78 @@ +"""LDS against a precomputed subset-retrain ground truth. + +``lds_from_precomputed_subsets`` is the scoring side of a published bank +(kronfluence's ``masks.pt``/``losses.pt``, say): the prediction for a subset is +the summed score of its members, and the LDS is the per-query Spearman against +the measured subset losses. These cover the group-prediction arithmetic, the +per-query averaging, and the shape guards that catch a scores/masks mismatch. +""" + +import numpy as np +import pytest +import torch + +from bergson.validate import lds_from_precomputed_subsets + + +@pytest.fixture +def bank(tmp_path): + """A 6-subset bank over 5 training rows and 2 queries, with losses made + exactly linear in the summed scores so the LDS is 1.0 by construction.""" + rng = np.random.default_rng(0) + masks = rng.integers(0, 2, size=(6, 5)).astype(np.float64) + scores = rng.normal(size=(5, 2)) + + torch.save(torch.from_numpy(masks), tmp_path / "masks.pt") + torch.save(torch.from_numpy(masks @ scores), tmp_path / "losses.pt") + np.save(tmp_path / "scores.npy", scores.astype(np.float32)) + return tmp_path + + +def test_perfectly_predictive_scores_score_one(bank): + lds = lds_from_precomputed_subsets( + str(bank / "scores.npy"), str(bank / "masks.pt"), str(bank / "losses.pt") + ) + assert lds == pytest.approx(1.0) + + +def test_sign_flipped_scores_score_minus_one(bank): + scores = -np.load(bank / "scores.npy") + np.save(bank / "flipped.npy", scores) + + lds = lds_from_precomputed_subsets( + str(bank / "flipped.npy"), str(bank / "masks.pt"), str(bank / "losses.pt") + ) + assert lds == pytest.approx(-1.0) + + +def test_summary_csv_has_one_row_per_query_plus_mean(bank): + out = bank / "lds.csv" + lds = lds_from_precomputed_subsets( + str(bank / "scores.npy"), + str(bank / "masks.pt"), + str(bank / "losses.pt"), + summary_path=str(out), + ) + + rows = out.read_text().strip().splitlines() + assert len(rows) == 1 + 2 + 1 # header, two queries, mean + assert rows[-1].startswith("mean,") + assert float(rows[-1].split(",")[1]) == pytest.approx(lds) + + +def test_mismatched_training_row_count_is_rejected(bank): + np.save(bank / "short.npy", np.zeros((4, 2), dtype=np.float32)) + + with pytest.raises(ValueError, match="masks cover 5 training rows"): + lds_from_precomputed_subsets( + str(bank / "short.npy"), str(bank / "masks.pt"), str(bank / "losses.pt") + ) + + +def test_mismatched_query_count_is_rejected(bank): + np.save(bank / "wide.npy", np.zeros((5, 3), dtype=np.float32)) + + with pytest.raises(ValueError, match="losses shape"): + lds_from_precomputed_subsets( + str(bank / "wide.npy"), str(bank / "masks.pt"), str(bank / "losses.pt") + ) diff --git a/tests/test_source_fisher_normalization.py b/tests/test_source_fisher_normalization.py new file mode 100644 index 00000000..559f0771 --- /dev/null +++ b/tests/test_source_fisher_normalization.py @@ -0,0 +1,83 @@ +"""Normalization of the segment EK-FAC eigenvalues fed to SOURCE's eigenfunctions. + +kronfluence divides its lambda matrix by a per-sample count before use +(``kronfluence/factor/config.py`` ``lambda_matrix.div_(NUM_LAMBDA_PROCESSED)``, +counted in ``module/tracker/factor.py`` as ``per_sample_gradient.size(0)``). +``LambdaCollector`` accumulates an unnormalized sum, so the segment aggregation +divides by the pooled count to match. +""" + +import json + +import pytest +import torch +from safetensors.torch import load_file, save_file + +from bergson.approx_unrolling.segment_aggregation import ( + LAMBDA_COUNTS_FILENAME, + _sum_my_shard, + lambda_denominator, + write_lambda_counts, +) + +DOCS, TOKENS = 4599, 2354682 + + +def _ckpt_dirs(tmp_path, n=2): + dirs = [] + for i in range(n): + d = tmp_path / f"ckpt_{i}" + d.mkdir() + write_lambda_counts(d, documents=DOCS, tokens=TOKENS) + dirs.append(d) + return dirs + + +def test_counts_roundtrip(tmp_path): + write_lambda_counts(tmp_path, documents=DOCS, tokens=TOKENS) + with open(tmp_path / LAMBDA_COUNTS_FILENAME) as f: + assert json.load(f) == {"documents": DOCS, "tokens": TOKENS} + + +@pytest.mark.parametrize( + "normalization, expected", + [("document", 2 * DOCS), ("token", 2 * TOKENS), ("none", 1.0)], +) +def test_denominator_pools_over_checkpoints(tmp_path, normalization, expected): + """The lambdas are summed over checkpoints, so the denominator sums too.""" + assert lambda_denominator(_ckpt_dirs(tmp_path), normalization) == expected + + +def test_denominator_reports_missing_counts(tmp_path): + d = tmp_path / "ckpt_0" + d.mkdir() + with pytest.raises(FileNotFoundError, match="fisher_normalization"): + lambda_denominator([d], "document") + + +def test_denominator_ignores_missing_counts_when_unnormalized(tmp_path): + """ "none" reproduces pre-normalization runs, which have no counts.json.""" + d = tmp_path / "ckpt_0" + d.mkdir() + assert lambda_denominator([d], "none") == 1.0 + + +def test_sum_my_shard_divides(tmp_path): + ins = [] + for i in range(2): + p = tmp_path / f"in_{i}.safetensors" + save_file({"w": torch.full((2, 3), float(i + 1))}, p) + ins.append(p) + + out = tmp_path / "out.safetensors" + _sum_my_shard(ins, out, device="cpu", divisor=6.0) + # (1 + 2) / 6 + assert torch.allclose(load_file(out)["w"], torch.full((2, 3), 0.5)) + + +def test_sum_my_shard_defaults_to_plain_sum(tmp_path): + p = tmp_path / "in.safetensors" + save_file({"w": torch.ones(2, 3)}, p) + out = tmp_path / "out.safetensors" + _sum_my_shard([p], out, device="cpu") + assert torch.allclose(load_file(out)["w"], torch.ones(2, 3)) diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index f562fef2..86292c69 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -22,13 +22,15 @@ write_lr_history, ) from bergson.config.config import ApproxUnrollingConfig, TrainingConfig +from bergson.config.config_io import save_run_config def _run_dir(tmp_path, **overrides): - """A bergson run directory with a config.yaml, as save_run_config writes it.""" - cfg = TrainingConfig(run_path=str(tmp_path), **overrides) - payload = {"magic": cfg.to_dict()} - (tmp_path / "config.yaml").write_text(yaml.safe_dump([payload])) + """A bergson run directory with a config.yaml. + + Written by ``save_run_config`` itself so the loader is pinned against the + real on-disk shape rather than a hand-rolled approximation of it.""" + save_run_config(TrainingConfig(run_path=str(tmp_path), **overrides), tmp_path) return tmp_path @@ -137,6 +139,23 @@ def test_missing_config_yaml_is_explained(tmp_path): load_training_config(tmp_path) +def test_load_reads_the_saved_steps_document(tmp_path): + """save_run_config wraps the step list in {steps, metadata}.""" + _run_dir(tmp_path, optimizer="adamw", adam_beta2=0.98) + + doc = yaml.safe_load((tmp_path / "config.yaml").read_text()) + assert set(doc) == {"steps", "metadata"} + assert load_training_config(tmp_path).adam_beta2 == pytest.approx(0.98) + + +def test_load_still_reads_a_bare_step_list(tmp_path): + """Runs saved before the {steps, metadata} wrapper stay readable.""" + cfg = TrainingConfig(run_path=str(tmp_path), optimizer="sgd", adam_beta1=0.8) + (tmp_path / "config.yaml").write_text(yaml.safe_dump([{"magic": cfg.to_dict()}])) + + assert load_training_config(tmp_path).adam_beta1 == pytest.approx(0.8) + + # ── LR history ────────────────────────────────────────────────────────────── @@ -504,3 +523,20 @@ def test_export_checkpoints_end_to_end(tmp_path): out = resolve(ApproxUnrollingConfig(checkpoints=[str(p) for p in exported])) assert out.model_path == model_name assert out.momentum == 0.0 # adamw + + +def test_resolve_ignores_a_non_training_config(tmp_path): + """A checkpoint dir's sibling config.yaml may belong to an attribution run. + + ``infer_trainer_run`` only checks that a config.yaml exists, so ``resolve`` + has to tolerate one it cannot read as a TrainingConfig. + """ + run = tmp_path / "attribution_run" + (run / "models" / "step_1").mkdir(parents=True) + (run / "config.yaml").write_text("steps:\n- approxunrolling:\n index_cfg: {}\n") + + cfg = ApproxUnrollingConfig(checkpoints=[str(run / "models" / "step_1")]) + resolved = resolve(cfg) + + assert resolved.momentum == 0.0 + assert resolved.model_path is None