diff --git a/bergson/approx_unrolling/adam_preconditioner.py b/bergson/approx_unrolling/adam_preconditioner.py new file mode 100644 index 00000000..05e9902e --- /dev/null +++ b/bergson/approx_unrolling/adam_preconditioner.py @@ -0,0 +1,180 @@ +"""Per-segment Adam SOURCE preconditioners. + +:func:`build_segment_preconditioners` reads each checkpoint dir's +``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``. +""" + +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, +) +from ..utils.trainer_export import OPTIMIZER_STATE_FILE + +__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 + 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" + 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, 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.""" + 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]: + """Get (beta2, eps, eps_root) from a standard ``optimizer.pt``. + + ``betas`` and ``eps`` are standard AdamW param-group keys (HF Trainer + 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; save " + "with TrainingConfig.save_optimizer_state " + ) + 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 an + ``optimizer.pt`` with 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; save " + "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``. 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 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) + 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)) + preconditioner = {} + 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) + preconditioner[module] = _orient(p, grids[module], module, layer) + + out_path = base / f"segment_{seg}" / "preconditioner.safetensors" + 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 bf093272..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,25 +129,51 @@ 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. - - ``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. - """ + """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``. + + 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).""" cfg = EkfacConfig( hessian_method_path=str(segment_dir), gradient_path=str(src_grad_path), run_path=str(dst_grad_path), ev_correction=True, + preconditioner_path=preconditioner_path, + # 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", @@ -163,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() @@ -178,6 +208,9 @@ def walk_query_phase1( method: str, 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. @@ -203,6 +236,9 @@ def walk_query_phase1( lr_times_steps=lr_times_steps_per_segment[k], 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 @@ -215,6 +251,9 @@ def walk_query_phase2( lr_times_steps_per_segment: list[float], 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. @@ -238,6 +277,9 @@ def walk_query_phase2( lr_times_steps=lr_times_steps_per_segment[l], 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) @@ -247,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." ) - 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 a2853c8b..88f8e872 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 save_optimizer_state='all', then " + "export with bergson.utils.trainer_export.export_checkpoints." + ) + lr_times_steps_per_segment = compute_lr_times_steps_per_segment( approx_unrolling_cfg ) @@ -126,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. @@ -170,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 @@ -197,6 +213,24 @@ def approx_unrolling_pipeline( ) build(query_cfg, query_preprocess_cfg) + # ── Per-segment Adam preconditioners from the checkpoints' second moments + 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: + 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: + preconditioner_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 +238,23 @@ 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 preconditioner_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 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 @@ -223,6 +269,9 @@ 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 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 @@ -232,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/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 b11a117e..5a4333fc 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -379,6 +379,11 @@ class TrainingConfig(AttributionConfig, Serializable): wandb_project: str = "" """Weights & Biases project name. If set, logs training loss to W&B.""" + save_models: bool = False + """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__() @@ -428,11 +433,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 +796,32 @@ 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`.""" + + 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 1f6f95a9..f9a334fa 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,27 @@ 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``.""" + 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``.""" + 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 + eigenfunction. Leave False for the backward eigenfunction, whose + P^1/2 exp(-t M) P^-1/2 sandwich cancels.""" debug: bool = False @@ -84,14 +105,49 @@ 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, - ) + chain: list = [] + if self.cfg.preconditioner_path: + if self.apply_fn is None: + raise ValueError("preconditioner_path requires apply_fn.") + diagonal = DiagonalFactoredPreconditioner.from_shards( + self.path, + self.cfg.preconditioner_path, + rank=self.rank, + device=self.device, + apply_fn=self.apply_fn, + 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, + 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] @@ -145,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/hessians/preconditioner.py b/bergson/hessians/preconditioner.py index cffb619b..ccaffbd7 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,173 @@ def apply(self, grads: dict[str, Tensor]) -> dict[str, Tensor]: return out +class DiagonalFactoredPreconditioner: + """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__( + self, + eigen_a: dict[str, Tensor], + eigen_g: dict[str, Tensor], + lambdas: dict[str, Tensor], + preconditioner: dict[str, Tensor], + *, + apply_fn, + multiply_by_preconditioner: bool = False, + ): + self.eigen_a = eigen_a + self.eigen_g = eigen_g + self.lambdas = lambdas + self.preconditioner = preconditioner + self.apply_fn = apply_fn + self.multiply_by_preconditioner = multiply_by_preconditioner + self.shard_computer = ShardedMul() + self.logger = get_logger("DiagonalFactoredPreconditioner") + self._multipliers: dict[str, Tensor] = {} + + @classmethod + def from_shards( + cls, + hessian_path: str | Path, + preconditioner_path: str | Path, + *, + rank: int, + device: str | torch.device, + apply_fn, + 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 ``preconditioner_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_grids = load_file(str(preconditioner_path), device=str(device)) + sharder = ShardedMul() + preconditioner = {} + for name, q_g in eigen_g.items(): + if name not in full_grids: + raise KeyError( + f"Module {name!r} has EKFAC factors but no preconditioner " + f"grid in {preconditioner_path}; available: " + f"{sorted(full_grids.keys())}" + ) + 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( + f"Preconditioner grid for {name!r} has shape " + f"{tuple(p.shape)}, expected [out, in] = ({o}, {i})." + ) + start, end = sharder.shard_bounds(o) + preconditioner[name] = p[start:end] + return cls( + eigen_a, + eigen_g, + lambdas, + preconditioner, + apply_fn=apply_fn, + multiply_by_preconditioner=multiply_by_preconditioner, + ) + + 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.preconditioner[name] + sigma = p * self._diag_hessian_shard(name) + multiplier = self.apply_fn(sigma) + if self.multiply_by_preconditioner: + 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] + # 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) + 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..d25eb87c 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -262,15 +262,18 @@ def worker( if getattr(run_cfg, "save_optimizer_state", "none") == "all" else None ), + save_interval=getattr(run_cfg, "save_interval", 0), ) - 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"), ) - if getattr(run_cfg, "save_retrained_models", False) 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. @@ -321,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") @@ -372,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 5621b266..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: @@ -115,7 +171,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 +386,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 +520,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 +530,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/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/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_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 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_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_optimizer_variants.py b/tests/test_source_optimizer_variants.py new file mode 100644 index 00000000..5309e45f --- /dev/null +++ b/tests/test_source_optimizer_variants.py @@ -0,0 +1,376 @@ +"""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 numpy as np +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.data import create_index, load_module_gradients +from bergson.hessians.apply_hessian import EkfacApplicator, EkfacConfig +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")) + 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): + """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 + 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, + preconditioner_path, + rank=0, + device="cpu", + apply_fn=fn, + multiply_by_preconditioner=fn_kind == "f_segment", + ev_correction=True, + ) + + grads = {MODULE: torch.randn(3, OUT_DIM * IN_DIM)} + 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": + 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 + preconditioner_path = _write_factor_shards(tmp_path, q_a, q_g, lam, precond) + + preconditioner = DiagonalFactoredPreconditioner.from_shards( + tmp_path, + preconditioner_path, + rank=0, + device="cpu", + apply_fn=f_segment(LR_TIMES_STEPS), + multiply_by_preconditioner=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_preconditioner_shape_mismatch_raises(tmp_path): + q_a, q_g, lam = _random_factors() + bad_precond = torch.rand(OUT_DIM + 1, IN_DIM) + 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, + preconditioner_path, + rank=0, + device="cpu", + apply_fn=f_backward(LR_TIMES_STEPS), + ev_correction=True, + ) + + +def test_preconditioner_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")) + 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, + preconditioner_path, + rank=0, + device="cpu", + apply_fn=f_backward(LR_TIMES_STEPS), + ev_correction=True, + ) + + +@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, + 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]) 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