From 5f8768f0494678630ffd91a31a92cc67467f62c6 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Fri, 31 Jul 2026 21:29:48 +0000 Subject: [PATCH 01/11] 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. --- .../approx_unrolling/approx_unrolling_math.py | 17 +-- bergson/approx_unrolling/pipeline.py | 3 +- .../precompute_checkpoints.py | 65 ++++++++--- .../approx_unrolling/segment_aggregation.py | 71 +++++++++++- bergson/approx_unrolling/trainer_run.py | 12 ++- bergson/config/config.py | 11 ++ ...pt2_wikitext_lotus_source_q50_docnorm.yaml | 83 ++++++++++++++ ...ext_lotus_source_q50_docnorm_validate.yaml | 18 ++++ ...pt2_wikitext_lotus_source_q50_toknorm.yaml | 83 ++++++++++++++ .../chain_toknorm.sh | 50 +++++++++ .../compare_lds.py | 81 ++++++++++++++ .../source_fisher_normalization/setup_run.py | 101 ++++++++++++++++++ tests/test_source_fisher_normalization.py | 83 ++++++++++++++ tests/test_source_trainer_integration.py | 17 +++ 14 files changed, 663 insertions(+), 32 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 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 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..3ea972a3 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 @@ -134,16 +134,17 @@ def apply_eigfn_to_query( 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 = "", ) -> 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).""" @@ -169,7 +170,7 @@ 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) diff --git a/bergson/approx_unrolling/pipeline.py b/bergson/approx_unrolling/pipeline.py index 971ea337..9a271aae 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 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..418074a9 100644 --- a/bergson/approx_unrolling/trainer_run.py +++ b/bergson/approx_unrolling/trainer_run.py @@ -127,7 +127,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..68858a74 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -802,6 +802,17 @@ class ApproxUnrollingConfig(Serializable): Bae et al. 2024, Appendix C). Requires an ``optimizer.pt`` in every checkpoint dir; see :mod:`bergson.approx_unrolling.adam_preconditioner`.""" + 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/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/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..5aa8961d --- /dev/null +++ b/experiments/source_fisher_normalization/setup_run.py @@ -0,0 +1,101 @@ +"""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/tests/test_source_fisher_normalization.py b/tests/test_source_fisher_normalization.py new file mode 100644 index 00000000..4e0b2e81 --- /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..882421b7 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -504,3 +504,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 From 4a5d408278e1daf3f0707aa0c0c25cb027a448a4 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Tue, 21 Jul 2026 08:59:51 +0000 Subject: [PATCH 02/11] feat: per-snapshot optimizer.pt export and interval save mode in MAGIC trainer (cherry picked from commit 16c29730ff6fdadd1c61ea72dc9dda43e6a85274) --- .../approx_unrolling/approx_unrolling_math.py | 85 ++++++++++-------- bergson/approx_unrolling/pipeline.py | 12 ++- bergson/magic/cli.py | 1 + bergson/magic/config.py | 9 +- bergson/magic/metasmoothness.py | 87 +++++++++++++++++++ bergson/magic/trainer.py | 27 +++++- 6 files changed, 181 insertions(+), 40 deletions(-) diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index 3ea972a3..95d544d8 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -258,51 +258,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." ) - if seg_score_cfg.higher_is_better: - seg_scores = -seg_scores + 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, + ) + 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 9a271aae..198de884 100644 --- a/bergson/approx_unrolling/pipeline.py +++ b/bergson/approx_unrolling/pipeline.py @@ -277,7 +277,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/magic/cli.py b/bergson/magic/cli.py index e6a134da..0bafd3c2 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. diff --git a/bergson/magic/config.py b/bergson/magic/config.py index 2fa0f944..dd3a07f8 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.""" 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..8a8e5b4e 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,21 @@ 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. + if save_dir 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): From 4bfe6720c7afa4587094d342c39b6575e42aefe7 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Fri, 31 Jul 2026 14:26:32 +0000 Subject: [PATCH 03/11] 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). --- bergson/magic/trainer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/bergson/magic/trainer.py b/bergson/magic/trainer.py index 8a8e5b4e..9acf7015 100644 --- a/bergson/magic/trainer.py +++ b/bergson/magic/trainer.py @@ -663,7 +663,9 @@ def train( # 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. - if save_dir and next_save == 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) From 1c143536b921fe44e758b1c3120efa99e940a014 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 22 Jul 2026 06:46:41 +0000 Subject: [PATCH 04/11] 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) --- bergson/magic/cli.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index 0bafd3c2..f009d6a0 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -324,7 +324,7 @@ def worker( ) multi_query = False - if not score_path: + if do_metagradient: # Sanity check if not isinstance(run_cfg, MagicConfig): raise RuntimeError("run_cfg must be a MagicConfig to compute scores") @@ -375,6 +375,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: From 875d9da2c2237510c2abc31eda3eca14b3ffcfd6 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Fri, 31 Jul 2026 14:29:30 +0000 Subject: [PATCH 05/11] 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. --- bergson/magic/cli.py | 5 +++++ bergson/magic/config.py | 6 ++++++ 2 files changed, 11 insertions(+) diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index f009d6a0..d25eb87c 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -324,6 +324,11 @@ def worker( ) multi_query = False + # 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): diff --git a/bergson/magic/config.py b/bergson/magic/config.py index dd3a07f8..b95a0bdd 100644 --- a/bergson/magic/config.py +++ b/bergson/magic/config.py @@ -43,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.""" From 17a900fd6d03d9a1ace381b78ac9a1a7b226a90f Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Fri, 31 Jul 2026 14:33:35 +0000 Subject: [PATCH 06/11] 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. --- .../approx_unrolling/approx_unrolling_math.py | 37 ++++++++++++++++++- bergson/approx_unrolling/pipeline.py | 4 ++ bergson/config/config.py | 9 +++++ bergson/hessians/apply_hessian.py | 32 +++++++++++++++- 4 files changed, 79 insertions(+), 3 deletions(-) diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index 95d544d8..669f880c 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -129,6 +129,19 @@ 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, @@ -137,6 +150,8 @@ def apply_eigfn_to_query( fn_kind: Literal["f_segment", "f_backward"], distributed: DistributedConfig, preconditioner_path: str = "", + preconditioner_hybrid: bool = False, + preconditioner_hybrid_damping: float = 0.1, ) -> None: """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 @@ -154,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", @@ -176,7 +195,13 @@ def _apply_eigfn_worker( # 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() @@ -186,6 +211,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. @@ -212,6 +239,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 @@ -225,6 +254,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. @@ -249,6 +280,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) diff --git a/bergson/approx_unrolling/pipeline.py b/bergson/approx_unrolling/pipeline.py index 198de884..88f8e872 100644 --- a/bergson/approx_unrolling/pipeline.py +++ b/bergson/approx_unrolling/pipeline.py @@ -253,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 @@ -268,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 diff --git a/bergson/config/config.py b/bergson/config/config.py index 68858a74..5a4333fc 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -802,6 +802,15 @@ 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. 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 From 4f44f4d0a8e1e9f2382fd604dc66cd54a8bcb128 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Fri, 31 Jul 2026 14:35:50 +0000 Subject: [PATCH 07/11] 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. --- examples/replication/wikitext_gpt2_ekfac.yaml | 56 +++++++++++++++ .../replication/wikitext_gpt2_source.yaml | 71 +++++++++++++++++++ examples/replication/wikitext_gpt2_train.yaml | 67 +++++++++++++++++ 3 files changed, 194 insertions(+) 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 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..96e8697b --- /dev/null +++ b/examples/replication/wikitext_gpt2_source.yaml @@ -0,0 +1,71 @@ +# 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). +# +# 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/models/step_291 + - runs/wikitext_repro/train/models/step_582 + - runs/wikitext_repro/train/models/step_873 + - runs/wikitext_repro/train/models/step_1164 + - runs/wikitext_repro/train/models/step_1455 + - runs/wikitext_repro/train/models/step_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 From 7c0b687528436aecab5bdf72a0e12780762abfc8 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Fri, 31 Jul 2026 14:41:57 +0000 Subject: [PATCH 08/11] 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. --- bergson/validate.py | 59 ++++++++++++++ .../replication/wikitext_gpt2_source.yaml | 17 ++-- tests/test_lds_precomputed_subsets.py | 78 +++++++++++++++++++ 3 files changed, 148 insertions(+), 6 deletions(-) create mode 100644 tests/test_lds_precomputed_subsets.py diff --git a/bergson/validate.py b/bergson/validate.py index 447613d2..4f7c9090 100644 --- a/bergson/validate.py +++ b/bergson/validate.py @@ -80,6 +80,65 @@ 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/replication/wikitext_gpt2_source.yaml b/examples/replication/wikitext_gpt2_source.yaml index 96e8697b..b118dd47 100644 --- a/examples/replication/wikitext_gpt2_source.yaml +++ b/examples/replication/wikitext_gpt2_source.yaml @@ -19,6 +19,11 @@ # 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 @@ -47,12 +52,12 @@ steps: approx_unrolling_cfg: checkpoints: - - runs/wikitext_repro/train/models/step_291 - - runs/wikitext_repro/train/models/step_582 - - runs/wikitext_repro/train/models/step_873 - - runs/wikitext_repro/train/models/step_1164 - - runs/wikitext_repro/train/models/step_1455 - - runs/wikitext_repro/train/models/step_1746 + - 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] 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") + ) From cf91c5b67892ff4e4c830ba816099d5335033962 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Fri, 31 Jul 2026 14:57:35 +0000 Subject: [PATCH 09/11] 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. --- bergson/approx_unrolling/trainer_run.py | 19 ++++++++++++----- tests/test_source_trainer_integration.py | 27 ++++++++++++++++++++---- 2 files changed, 37 insertions(+), 9 deletions(-) diff --git a/bergson/approx_unrolling/trainer_run.py b/bergson/approx_unrolling/trainer_run.py index 418074a9..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: diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index 882421b7..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 ────────────────────────────────────────────────────────────── From 1a8d9d1cb1ab8c78dc0886a34f5f6b97fee09e7f Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Fri, 31 Jul 2026 21:29:40 +0000 Subject: [PATCH 10/11] 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. --- experiments/wikitext_replication/run_chain.py | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 experiments/wikitext_replication/run_chain.py 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() From 651a57550d1e15373c346bdc3b9ad3077131e0a0 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Fri, 31 Jul 2026 21:30:33 +0000 Subject: [PATCH 11/11] style: apply black formatting --- bergson/approx_unrolling/approx_unrolling_math.py | 4 +--- bergson/validate.py | 5 +---- experiments/source_fisher_normalization/setup_run.py | 10 ++++++++-- tests/test_source_fisher_normalization.py | 2 +- 4 files changed, 11 insertions(+), 10 deletions(-) diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index 669f880c..ddc886fc 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -199,9 +199,7 @@ def _apply_eigfn_worker( # 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 - ) + fn = {"f_segment": f_segment, "f_backward": f_backward}[fn_kind](lr_times_steps) EkfacApplicator(cfg, apply_fn=fn).compute_ivhp_sharded() diff --git a/bergson/validate.py b/bergson/validate.py index 4f7c9090..ba824707 100644 --- a/bergson/validate.py +++ b/bergson/validate.py @@ -122,10 +122,7 @@ def lds_from_precomputed_subsets( preds = masks @ scores.double().numpy() per_query = np.array( - [ - spearmanr(preds[:, q], losses[:, q]).statistic - for q in range(losses.shape[1]) - ] + [spearmanr(preds[:, q], losses[:, q]).statistic for q in range(losses.shape[1])] ) mean_lds = float(np.nanmean(per_query)) diff --git a/experiments/source_fisher_normalization/setup_run.py b/experiments/source_fisher_normalization/setup_run.py index 5aa8961d..62efa7f5 100644 --- a/experiments/source_fisher_normalization/setup_run.py +++ b/experiments/source_fisher_normalization/setup_run.py @@ -85,7 +85,11 @@ def main() -> None: 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",) + 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(): @@ -94,7 +98,9 @@ def main() -> None: ) print(f"hardlinked {total} files from {baseline} -> {target}") - print("not linked (regenerated): eigenvalue_correction_sharded, query_grad_*, scores") + print( + "not linked (regenerated): eigenvalue_correction_sharded, query_grad_*, scores" + ) if __name__ == "__main__": diff --git a/tests/test_source_fisher_normalization.py b/tests/test_source_fisher_normalization.py index 4e0b2e81..559f0771 100644 --- a/tests/test_source_fisher_normalization.py +++ b/tests/test_source_fisher_normalization.py @@ -56,7 +56,7 @@ def test_denominator_reports_missing_counts(tmp_path): def test_denominator_ignores_missing_counts_when_unnormalized(tmp_path): - """"none" reproduces pre-normalization runs, which have no counts.json.""" + """ "none" reproduces pre-normalization runs, which have no counts.json.""" d = tmp_path / "ckpt_0" d.mkdir() assert lambda_denominator([d], "none") == 1.0