diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index 4ec7b518..ddc886fc 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -3,7 +3,7 @@ import shutil from copy import deepcopy from pathlib import Path -from typing import Callable +from typing import Callable, Literal import numpy as np import torch @@ -129,21 +129,37 @@ def fn(sigma: Tensor) -> Tensor: return fn +def f_one_minus_exp(lr_times_steps: float) -> Callable[[Tensor], Tensor]: + """x -> 1 - exp(-lr_times_steps*x), the numerator of :func:`f_segment`. + + Used by the Eq-43 hybrid: the 1/x that completes f_segment is supplied by + the chained EK-FAC inverse rather than evaluated on the diagonal. + """ + + def fn(sigma: Tensor) -> Tensor: + return -torch.expm1(-lr_times_steps * sigma) + + return fn + + def apply_eigfn_to_query( src_grad_path: Path, dst_grad_path: Path, segment_dir: Path, lr_times_steps: float, - fn_kind: str, + fn_kind: Literal["f_segment", "f_backward"], distributed: DistributedConfig, preconditioner_path: str = "", + preconditioner_hybrid: bool = False, + preconditioner_hybrid_damping: float = 0.1, ) -> None: - """Apply F_segment or F_backward of one segment to a stored query gradient. + """Apply a segment's f_segment or f_backward operator to a stored query + gradient. The operator scales each query element in the segment's + eigenbasis by the function's value at the matching eigenvalue, averaged + over the segment's (checkpoint, document) pairs by + ``ApproxUnrollingConfig.fisher_normalization``. - ``fn_kind`` is "f_segment" or "f_backward". The segment eigenvalues are - already checkpoint-averaged (expected eigenvalues), so the eigenfunction is - applied to them directly. - ``preconditioner_path`` selects the preconditioned-optimizer variant: the + When ``preconditioner_path`` is set for an Adam-optimized run the eigenfunction is evaluated on a diagonal approximation of P^1/2 H P^1/2 in parameter space, with F_segment's output additionally multiplied by P (its P^1/2 . P^1/2 sandwich; F_backward's sandwich cancels).""" @@ -153,7 +169,11 @@ def apply_eigfn_to_query( run_path=str(dst_grad_path), ev_correction=True, preconditioner_path=preconditioner_path, - preconditioner_post_multiply=fn_kind == "f_segment", + # The hybrid's P^1/2 sandwich is absorbed by the chained EK-FAC inverse. + preconditioner_post_multiply=fn_kind == "f_segment" + and not preconditioner_hybrid, + preconditioner_hybrid=preconditioner_hybrid and fn_kind == "f_segment", + preconditioner_hybrid_damping=preconditioner_hybrid_damping, ) launch_distributed_run( "apply_eigfn_to_query", @@ -169,13 +189,17 @@ def _apply_eigfn_worker( world_size: int, cfg: EkfacConfig, lr_times_steps: float, - fn_kind: str, + fn_kind: Literal["f_segment", "f_backward"], ) -> None: init_dist(rank, local_rank, world_size) # Segment eigenvalues are already checkpoint-averaged, so the eigenfunction # is applied to them directly (no per-example normalization). - fn = {"f_segment": f_segment, "f_backward": f_backward}[fn_kind](lr_times_steps) + if cfg.preconditioner_hybrid: + # 1/x comes from the chained EK-FAC inverse, so mask with the numerator. + fn = f_one_minus_exp(lr_times_steps) + else: + fn = {"f_segment": f_segment, "f_backward": f_backward}[fn_kind](lr_times_steps) EkfacApplicator(cfg, apply_fn=fn).compute_ivhp_sharded() @@ -185,6 +209,8 @@ def walk_query_phase1( lr_times_steps_per_segment: list[float], distributed: DistributedConfig, preconditioner_paths: list[str] | None = None, + preconditioner_hybrid: bool = False, + preconditioner_hybrid_damping: float = 0.1, ) -> list[Path]: """Phase 1: build query_grad_0, ..., query_grad_{L-1} by walking F_backward. @@ -211,6 +237,8 @@ def walk_query_phase1( fn_kind="f_backward", distributed=distributed, preconditioner_path=preconditioner_paths[k] if preconditioner_paths else "", + preconditioner_hybrid=preconditioner_hybrid, + preconditioner_hybrid_damping=preconditioner_hybrid_damping, ) query_grad_paths[k - 1] = dst @@ -224,6 +252,8 @@ def walk_query_phase2( query_grad_paths: list[Path], distributed: DistributedConfig, preconditioner_paths: list[str] | None = None, + preconditioner_hybrid: bool = False, + preconditioner_hybrid_damping: float = 0.1, ) -> list[Path]: """Phase 2: build query_grad_segment_0, ..., query_grad_segment_{L-1} via F_segment. @@ -248,6 +278,8 @@ def walk_query_phase2( fn_kind="f_segment", distributed=distributed, preconditioner_path=preconditioner_paths[l] if preconditioner_paths else "", + preconditioner_hybrid=preconditioner_hybrid, + preconditioner_hybrid_damping=preconditioner_hybrid_damping, ) query_grad_segment_paths.append(dst) @@ -257,51 +289,64 @@ def walk_query_phase2( def score_per_segment_and_aggregate( index_cfg: IndexConfig, query_grad_segment_paths: list[Path], - final_checkpoint: str, + segment_checkpoints: list[list[str]], query_batch_size: int | None = None, ) -> Path: - """Phase 3: per-segment ``query_grad_segment_l . g(z_m)`` scores, summed. - - For each l, runs :func:`score_dataset` against the training data at the - final checkpoint with ``query_grad_segment_l`` as the query. Writes - per-segment outputs to ``/segment_{l}/scores/``, then sums into - ``/scores/``. + """Phase 3: per-segment ``query_grad_segment_l . g_bar_l(z_m)`` scores, summed. + + ``g_bar_l`` is the segment's expected training gradient (Bae et al. 2024, + Sec. 3.4; App. D notes the scores need "the training gradients C times, + where C is the total number of checkpoints"). Per-example gradients are + computed at each checkpoint *within* segment l and averaged; scoring every + segment at the final checkpoint instead collapses SOURCE toward influence + functions at convergence. + + Scores are linear in the training gradient, so averaging the per-checkpoint + scores is equivalent to scoring against the averaged gradient. Per-checkpoint + outputs land at ``/segment_{l}/scores_ckpt_{c}/``. """ base_run = Path(index_cfg.run_path) - num_segments = len(query_grad_segment_paths) score_dirs: list[Path] = [] - for l in range(num_segments): - scores_dir = base_run / f"segment_{l}" / "scores" - if index_cfg.distributed._node_rank == 0 and scores_dir.exists(): - shutil.rmtree(scores_dir) - seg_index_cfg = deepcopy(index_cfg) - seg_index_cfg.model = final_checkpoint - seg_index_cfg.run_path = str(scores_dir) - seg_index_cfg.projection_dim = 0 - score_cfg = ScoreConfig( - query_path=str(query_grad_segment_paths[l]), - higher_is_better=True, - query_batch_size=query_batch_size, - ) - seg_preprocess_cfg = PreprocessConfig() - save_run_config( - Score(score_cfg, seg_index_cfg, seg_preprocess_cfg), - seg_index_cfg.partial_run_path, - ) - score_dataset(seg_index_cfg, score_cfg, seg_preprocess_cfg) - score_dirs.append(scores_dir) - total = None - for scores_dir in score_dirs: - seg_scores = load_scores(scores_dir)[:] - seg_score_cfg = load_subconfig(scores_dir, "score_cfg", ScoreConfig) - if seg_score_cfg is None: + def _oriented(scores_dir: Path): + scores = load_scores(scores_dir)[:] + score_cfg = load_subconfig(scores_dir, "score_cfg", ScoreConfig) + if score_cfg is None: raise FileNotFoundError( f"No score_cfg found at {scores_dir}; cannot determine the " - "segment scores' orientation for aggregation." + "scores' orientation for aggregation." + ) + return -scores if score_cfg.higher_is_better else scores + + total = None + for l, ckpts in enumerate(segment_checkpoints): + seg_total = None + for c, ckpt in enumerate(ckpts): + scores_dir = base_run / f"segment_{l}" / f"scores_ckpt_{c}" + if index_cfg.distributed._node_rank == 0 and scores_dir.exists(): + shutil.rmtree(scores_dir) + seg_index_cfg = deepcopy(index_cfg) + seg_index_cfg.model = ckpt + seg_index_cfg.run_path = str(scores_dir) + seg_index_cfg.projection_dim = 0 + score_cfg = ScoreConfig( + query_path=str(query_grad_segment_paths[l]), + higher_is_better=True, + query_batch_size=query_batch_size, + ) + seg_preprocess_cfg = PreprocessConfig() + save_run_config( + Score(score_cfg, seg_index_cfg, seg_preprocess_cfg), + seg_index_cfg.partial_run_path, ) - if seg_score_cfg.higher_is_better: - seg_scores = -seg_scores + score_dataset(seg_index_cfg, score_cfg, seg_preprocess_cfg) + score_dirs.append(scores_dir) + + ckpt_scores = _oriented(scores_dir) + seg_total = ckpt_scores if seg_total is None else seg_total + ckpt_scores + + assert seg_total is not None, "each segment has >= 1 checkpoint" + seg_scores = seg_total / len(ckpts) total = seg_scores if total is None else total + seg_scores assert total is not None, "num_segments >= 1 is validated by the pipeline" diff --git a/bergson/approx_unrolling/pipeline.py b/bergson/approx_unrolling/pipeline.py index 971ea337..88f8e872 100644 --- a/bergson/approx_unrolling/pipeline.py +++ b/bergson/approx_unrolling/pipeline.py @@ -141,7 +141,7 @@ def approx_unrolling_pipeline( index_cfg, hessian_cfg, approx_unrolling_cfg, - overwrite=index_cfg.overwrite, + resume=index_cfg.overwrite, ) # Encourage GC between expensive steps; matters most in single-GPU mode # where the worker ran in-process and may still hold model references. @@ -185,6 +185,7 @@ def approx_unrolling_pipeline( per_segment=n_ckpts // n_segments, distributed=index_cfg.distributed, resume=index_cfg.overwrite, + normalization=approx_unrolling_cfg.fisher_normalization, ) # ── Step 5: Mean query gradient at the final checkpoint @@ -252,6 +253,8 @@ def approx_unrolling_pipeline( lr_times_steps_per_segment=lr_times_steps_per_segment, distributed=index_cfg.distributed, preconditioner_paths=[str(p) for p in preconditioner_paths] or None, + preconditioner_hybrid=approx_unrolling_cfg.adam_segment_hybrid, + preconditioner_hybrid_damping=approx_unrolling_cfg.adam_hybrid_damping, ) # ── Step 7: Phase 2 -- Get per-ckpt queries from segment queries @@ -267,6 +270,8 @@ def approx_unrolling_pipeline( query_grad_paths=query_grad_paths, distributed=index_cfg.distributed, preconditioner_paths=[str(p) for p in preconditioner_paths] or None, + preconditioner_hybrid=approx_unrolling_cfg.adam_segment_hybrid, + preconditioner_hybrid_damping=approx_unrolling_cfg.adam_hybrid_damping, ) # ── Step 8: Phase 3 -- per-segment scoring + sum @@ -276,7 +281,17 @@ def approx_unrolling_pipeline( out_path = score_per_segment_and_aggregate( index_cfg=index_cfg, query_grad_segment_paths=query_grad_segment_paths, - final_checkpoint=str(approx_unrolling_cfg.checkpoints[-1]), + # Bae et al. Sec. 3.4: g_bar_l averages the training gradients over the + # checkpoints within segment l, not the final checkpoint alone. + segment_checkpoints=[ + [ + str(c) + for c in approx_unrolling_cfg.checkpoints[ + l * (n_ckpts // n_segments) : (l + 1) * (n_ckpts // n_segments) + ] + ] + for l in range(n_segments) + ], query_batch_size=approx_unrolling_cfg.query_batch_size, ) logger.info(f"[approximate unrolling pipeline] DONE. Final scores at {out_path}") diff --git a/bergson/approx_unrolling/precompute_checkpoints.py b/bergson/approx_unrolling/precompute_checkpoints.py index a61ef7a3..8db2d579 100644 --- a/bergson/approx_unrolling/precompute_checkpoints.py +++ b/bergson/approx_unrolling/precompute_checkpoints.py @@ -2,6 +2,7 @@ from copy import deepcopy from pathlib import Path +import torch from datasets import Dataset from bergson.collector.collector import ( @@ -24,6 +25,8 @@ setup_model_and_peft, ) +from .segment_aggregation import LAMBDA_COUNTS_FILENAME, write_lambda_counts + def precompute_checkpoint_averaged_lambdas( index_cfg: IndexConfig, @@ -49,7 +52,14 @@ def precompute_checkpoint_averaged_lambdas( eigen_path = base_run / f"segment_{seg}" / method out_path = ckpt_method_dir / output_subdir - if out_path.exists(): + # A run predating `fisher_normalization` has lambdas but no counts.json; + # recover the counts without redoing the data pass. + counts_only = ( + out_path.exists() + and resume + and not (out_path / LAMBDA_COUNTS_FILENAME).exists() + ) + if out_path.exists() and not counts_only: if resume: logger.info( f"[seg {seg} ckpt {idx_in_seg}] skip — exists at {out_path}" @@ -63,8 +73,13 @@ def precompute_checkpoint_averaged_lambdas( ) logger.info( - f"[seg {seg} ckpt {idx_in_seg}] computing {output_subdir} at " - f"model={ckpt!r} using eigvecs from {eigen_path}" + f"[seg {seg} ckpt {idx_in_seg}] " + + ( + f"backfilling {LAMBDA_COUNTS_FILENAME} for existing {output_subdir}" + if counts_only + else f"computing {output_subdir} at model={ckpt!r} " + f"using eigvecs from {eigen_path}" + ) ) ckpt_index_cfg = deepcopy(index_cfg) @@ -77,18 +92,31 @@ def precompute_checkpoint_averaged_lambdas( ds, _ = setup_data_pipeline(ckpt_index_cfg) - launch_distributed_run( - "checkpoint_averaged_lambda", - _lambda_worker, - [ - ckpt_index_cfg, - hessian_cfg, - ds, - ckpt_method_dir, - eigen_path, - output_subdir, - ], - ckpt_index_cfg.distributed, + if not counts_only: + launch_distributed_run( + "checkpoint_averaged_lambda", + _lambda_worker, + [ + ckpt_index_cfg, + hessian_cfg, + ds, + ckpt_method_dir, + eigen_path, + output_subdir, + ], + ckpt_index_cfg.distributed, + ) + + write_lambda_counts( + out_path, + documents=len(ds), + tokens=int( + torch.load( + ckpt_index_cfg.partial_run_path / "total_processed.pt", + map_location="cpu", + weights_only=False, + ) + ), ) @@ -151,10 +179,13 @@ def precompute_checkpoint_hessians( hessian_cfg: HessianConfig, approx_unrolling_cfg: ApproxUnrollingConfig, *, - overwrite: bool = False, + resume: bool = False, ) -> None: """Run :func:`approximate_hessians` once per checkpoint. + ``resume`` skips checkpoints whose output already exists, matching the + other steps of the pipeline. + The pipeline-level divisibility check (``n_ckpts % n_segments == 0``) runs in :func:`approx_unrolling_pipeline` before this is called. """ @@ -173,7 +204,7 @@ def precompute_checkpoint_hessians( out_path = ckpt_dir / method if out_path.exists(): - if not overwrite: + if resume: logger.info( f"[seg {seg} ckpt {idx_in_seg}] skip — exists at {out_path}" ) diff --git a/bergson/approx_unrolling/segment_aggregation.py b/bergson/approx_unrolling/segment_aggregation.py index fcfdbb76..eb0650cb 100644 --- a/bergson/approx_unrolling/segment_aggregation.py +++ b/bergson/approx_unrolling/segment_aggregation.py @@ -1,5 +1,7 @@ +import json import shutil from pathlib import Path +from typing import Literal import torch import torch.distributed as dist @@ -14,18 +16,62 @@ _SHARD_KINDS = ("activation_sharded", "gradient_sharded") +LAMBDA_COUNTS_FILENAME = "counts.json" + +FisherNormalization = Literal["document", "token", "none"] + + +def write_lambda_counts(lambda_dir: Path, *, documents: int, tokens: int) -> None: + """Record the denominators available for normalizing a checkpoint's lambdas. + + ``LambdaCollector`` accumulates a sum over documents; ``total_processed`` + counts tokens. Both are recorded so the choice is made at aggregation time. + """ + with open(lambda_dir / LAMBDA_COUNTS_FILENAME, "w") as f: + json.dump({"documents": documents, "tokens": tokens}, f) + + +def lambda_denominator( + input_dirs: list[Path], normalization: FisherNormalization +) -> float: + """Pooled count over a segment's checkpoints. + + The lambdas are summed over checkpoints, so the denominator sums too: the + result is the mean over every (checkpoint, document) pair, which is what + kronfluence's ``lambda_matrix / num_lambda_processed`` computes for a single + checkpoint. + """ + if normalization == "none": + return 1.0 + + key = "documents" if normalization == "document" else "tokens" + total = 0 + for d in input_dirs: + counts_path = d / LAMBDA_COUNTS_FILENAME + if not counts_path.exists(): + raise FileNotFoundError( + f"Missing {counts_path}, needed to normalize the segment " + f"eigenvalues by {key}. It is written alongside the lambda " + "shards; re-run the per-checkpoint lambda step, or set " + "`fisher_normalization: none` to keep the unnormalized sum." + ) + with open(counts_path) as f: + total += json.load(f)[key] + return float(total) + def sum_sharded_dirs( input_dirs: list[Path], output_dir: Path, distributed: DistributedConfig, + divisor: float = 1.0, ) -> None: """Sum per-rank shards across ``input_dirs`` into ``output_dir``.""" output_dir.mkdir(parents=True, exist_ok=True) launch_distributed_run( "sum_sharded_dirs", _sum_sharded_dirs_worker, - [input_dirs, output_dir], + [input_dirs, output_dir, divisor], distributed, ) @@ -36,12 +82,13 @@ def _sum_sharded_dirs_worker( world_size: int, input_dirs: list[Path], output_dir: Path, + divisor: float, ) -> None: init_dist(rank, local_rank, world_size) device = get_device(local_rank) shard_name = f"shard_{rank}.safetensors" in_paths = [d / shard_name for d in input_dirs] - _sum_my_shard(in_paths, output_dir / shard_name, device=device) + _sum_my_shard(in_paths, output_dir / shard_name, device=device, divisor=divisor) if world_size > 1: dist.barrier() @@ -50,6 +97,7 @@ def _sum_my_shard( in_paths: list[Path], out_path: Path, device: str, + divisor: float = 1.0, ) -> None: """Sum all tensor dicts in ``in_paths`` and write to ``out_path``.""" acc: dict[str, torch.Tensor] = {} @@ -61,6 +109,9 @@ def _sum_my_shard( acc[k] = t.clone() else: acc[k].add_(t) + if divisor != 1.0: + for v in acc.values(): + v.div_(divisor) save_file({k: v.cpu() for k, v in acc.items()}, out_path) @@ -191,10 +242,16 @@ def aggregate_segment_lambdas( distributed: DistributedConfig, *, resume: bool = False, + normalization: FisherNormalization = "document", input_subdir: str = "averaged_ev_correct_sharded", output_subdir: str = "eigenvalue_correction_sharded", ) -> None: - """Sum per-checkpoint lambdas into per-segment lambda.""" + """Sum per-checkpoint lambdas into per-segment lambda. + + ``normalization`` divides the sum by the pooled document or token count, so + the segment eigenvalues are an expected Fisher rather than a total that + grows with dataset size. See ``ApproxUnrollingConfig.fisher_normalization``. + """ logger = get_logger("aggregate_segment_lambdas") base_run = Path(run_path) @@ -217,5 +274,9 @@ def aggregate_segment_lambdas( f"Missing per-ckpt lambda dir {d}; did step 3 finish?" ) - logger.info(f"[seg {seg}] summing lambdas -> {out_dir}") - sum_sharded_dirs(input_dirs, out_dir, distributed) + divisor = lambda_denominator(input_dirs, normalization) + logger.info( + f"[seg {seg}] summing lambdas -> {out_dir} " + f"(normalization={normalization}, divisor={divisor:g})" + ) + sum_sharded_dirs(input_dirs, out_dir, distributed, divisor=divisor) diff --git a/bergson/approx_unrolling/trainer_run.py b/bergson/approx_unrolling/trainer_run.py index 5e37dd9d..b5a8198e 100644 --- a/bergson/approx_unrolling/trainer_run.py +++ b/bergson/approx_unrolling/trainer_run.py @@ -40,8 +40,10 @@ def write_lr_history( def load_training_config(trainer_run: str | Path) -> TrainingConfig: - """Load the ``TrainingConfig`` a run was launched with. ``save_run_config`` - writes ``{command_name: {...}}``, so the single value is the config.""" + """Load the ``TrainingConfig`` a run was launched with. + + ``save_run_config`` writes ``{steps: [{command_name: {...}}], metadata: ...}``, + so the training config is the first step's single value.""" path = Path(trainer_run) / CONFIG_FILENAME if not path.is_file(): raise FileNotFoundError( @@ -50,9 +52,11 @@ def load_training_config(trainer_run: str | Path) -> TrainingConfig: ) with open(path) as f: - loaded = yaml.safe_load(f) + loaded: Any = yaml.safe_load(f) - # One-step configs are a list of {command: payload}; take the first payload. + if isinstance(loaded, dict) and "steps" in loaded: + loaded = loaded["steps"] + # Steps are a list of {command: payload}; take the first payload. if isinstance(loaded, list): if not loaded: raise ValueError(f"{path} is empty") @@ -64,7 +68,12 @@ def load_training_config(trainer_run: str | Path) -> TrainingConfig: if not isinstance(payload, dict): raise ValueError(f"{path} is not a bergson run config") - return TrainingConfig.from_dict(payload, drop_extra_fields=True) + try: + return TrainingConfig.from_dict(payload, drop_extra_fields=True) + except Exception as e: + # An attribution run's config.yaml has the same shape, so reaching here + # is expected; callers that guess at a run dir catch ValueError. + raise ValueError(f"{path} is not a bergson training config: {e}") from e def derive_momentum(training_cfg: TrainingConfig) -> float: @@ -127,7 +136,17 @@ def resolve(cfg: ApproxUnrollingConfig) -> ApproxUnrollingConfig: cfg.momentum = 0.0 return cfg - training_cfg = load_training_config(trainer_run) + try: + training_cfg = load_training_config(trainer_run) + except ValueError as e: + # A directory with a config.yaml is not necessarily a training run: an + # attribution run writes one next to the checkpoints it produced. Infer + # nothing from it rather than failing the pipeline. + logger.warning("Ignoring %s as a trainer run: %s", trainer_run, e) + if cfg.momentum is None: + cfg.momentum = 0.0 + return cfg + filled: list[str] = [] if cfg.model_path is None: diff --git a/bergson/config/config.py b/bergson/config/config.py index 0730f480..5a4333fc 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -802,6 +802,26 @@ class ApproxUnrollingConfig(Serializable): Bae et al. 2024, Appendix C). Requires an ``optimizer.pt`` in every checkpoint dir; see :mod:`bergson.approx_unrolling.adam_preconditioner`.""" + adam_segment_hybrid: bool = False + """With ``use_adam_preconditioner``: use the Eq-43 hybrid for F_segment -- + the diagonal supplies only the matrix exponential, the EK-FAC factors supply + H^-1 (Bae et al. 2024, App. D: "we still use the EK-FAC factors to compute + H^-1 g in Equation 43").""" + + adam_hybrid_damping: float = 0.1 + """Relative damping for the hybrid's EK-FAC inverse.""" + + fisher_normalization: Literal["document", "token", "none"] = "document" + """Denominator for the segment EK-FAC eigenvalues, which set the scale of + ``lr*steps * sigma`` that ``f_segment``/``f_backward`` are evaluated at. + + - "document": divide by the number of documents summed over the segment's + checkpoints. Matches kronfluence (Bae et al.'s implementation), which + divides its lambda matrix by a per-sample count. + - "token": divide by the token count instead. + - "none": leave the raw sum. Pre-normalization behaviour; the resulting + scale grows with dataset size and checkpoints per segment.""" + query: DataConfig = field(default_factory=DataConfig) """Query dataset spec; gradients computed at the final checkpoint.""" diff --git a/bergson/hessians/apply_hessian.py b/bergson/hessians/apply_hessian.py index ed17fb28..f9a334fa 100644 --- a/bergson/hessians/apply_hessian.py +++ b/bergson/hessians/apply_hessian.py @@ -45,6 +45,14 @@ class EkfacConfig: diagonal approximation of P^1/2 H P^1/2 in parameter space (the Adam/AdamW approximate-unrolling variant) instead of on the eigenvalues in the EKFAC eigenbasis. Requires ``apply_fn``.""" + preconditioner_hybrid: bool = False + """With ``preconditioner_path``: instead of evaluating ``apply_fn`` wholly on + the diagonal, apply the elementwise ``apply_fn(p * diag(H))`` mask and then + the standard EK-FAC inverse. This is the Eq-43 reading of the Adam SOURCE + variant (Bae et al. 2024, App. C/D), which keeps eigenbasis curvature in the + inverse rather than approximating it diagonally.""" + preconditioner_hybrid_damping: float = 0.1 + """Relative damping for the hybrid's EK-FAC inverse.""" preconditioner_post_multiply: bool = False """With ``preconditioner_path``: multiply the applied function's output by the preconditioner grid — the P^1/2 F(M) P^1/2 sandwich of the segment @@ -97,10 +105,11 @@ def __init__( self.device = get_device(self.rank) def compute_ivhp_sharded(self): + chain: list = [] if self.cfg.preconditioner_path: if self.apply_fn is None: raise ValueError("preconditioner_path requires apply_fn.") - preconditioner = DiagonalFactoredPreconditioner.from_shards( + diagonal = DiagonalFactoredPreconditioner.from_shards( self.path, self.cfg.preconditioner_path, rank=self.rank, @@ -109,6 +118,25 @@ def compute_ivhp_sharded(self): multiply_by_preconditioner=self.cfg.preconditioner_post_multiply, ev_correction=self.cfg.ev_correction, ) + if self.cfg.preconditioner_hybrid: + # Bae et al. App. D: "use the diagonal Hessian approximation for + # computing the matrix exponential ... Note that we still use the + # EK-FAC factors to compute H^-1 g in Equation 43." + # Eq-43 reads M_mask @ H^-1 @ g_train; applied to the QUERY + # gradient that is the adjoint H^-1 @ M_mask @ q, so the + # diagonal mask goes first and the EK-FAC inverse second. + chain.append(diagonal) + preconditioner = FactoredPreconditioner.from_shards( + self.path, + rank=self.rank, + device=self.device, + inversion_cfg=InversionConfig( + damping_factor=self.cfg.preconditioner_hybrid_damping + ), + ev_correction=self.cfg.ev_correction, + ) + else: + preconditioner = diagonal else: preconditioner = FactoredPreconditioner.from_shards( self.path, @@ -173,6 +201,8 @@ def compute_ivhp_sharded(self): device=self.device, dtype=torch.float32 ) + for pre in chain: + grads = pre.apply(grads) transformed = preconditioner.apply(grads) del grads diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index e6a134da..d25eb87c 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -262,6 +262,7 @@ def worker( if getattr(run_cfg, "save_optimizer_state", "none") == "all" else None ), + save_interval=getattr(run_cfg, "save_interval", 0), ) # Called on every rank: FSDP moments are DTensors whose gather is a # collective; rank 0 writes inside. @@ -323,7 +324,12 @@ def worker( ) multi_query = False - if not score_path: + # Plain runs compute scores here; skip_metagradient turns MAGIC into a + # trainer, which is what the SOURCE/EK-FAC replications want from it. + do_metagradient = not score_path and not getattr( + run_cfg, "skip_metagradient", False + ) + if do_metagradient: # Sanity check if not isinstance(run_cfg, MagicConfig): raise RuntimeError("run_cfg must be a MagicConfig to compute scores") @@ -374,6 +380,19 @@ def worker( score_path = os.path.join(run_cfg.run_path, "scores.pt") torch.save(scores, score_path) print(f"Saved attribution scores to {score_path}") + elif not score_path: + # skip_metagradient: no MAGIC scores; build the bank with dummy zeros. + scores = torch.zeros_like(stream.weights).cpu() + if pad_count: + scores = ( + scores[:-weight_pad_count] if scores.ndim == 1 else scores[:-pad_count] + ) + if global_rank == 0: + print(f"Baseline loss: {baseline}") + print( + "skip_metagradient: skipped MAGIC backward; using dummy zero " + "scores (bank build only)." + ) elif os.path.isdir(score_path) or score_path.endswith(".npy"): scores, multi_query = load_attribution_scores(score_path) else: diff --git a/bergson/magic/config.py b/bergson/magic/config.py index 2fa0f944..b95a0bdd 100644 --- a/bergson/magic/config.py +++ b/bergson/magic/config.py @@ -3,7 +3,7 @@ from ..config.config import ValidationConfig -MagicSaveMode = Literal["all", "sqrt", "log"] +MagicSaveMode = Literal["all", "sqrt", "log", "interval"] @dataclass @@ -20,10 +20,17 @@ class MagicConfig(ValidationConfig): - 'sqrt' saves at a linearly-spaced interval, every sqrt(N) steps. This method uses O(sqrt N) space and O(N) time. + - 'interval' saves every `save_interval` steps (plus the final state when the + cadence lands on it). Use for SOURCE-style runs that need a few evenly + spaced checkpoints rather than backward-replay coverage. + The original MAGIC paper used 'log', but 'sqrt' is often a better choice when disk space is not a concern. """ + save_interval: int = 0 + """Snapshot spacing in steps for `save_mode: interval`.""" + backward_save_every: int = 0 """How often (in steps) to save backward state for resume.""" @@ -36,3 +43,9 @@ class MagicConfig(ValidationConfig): skip_validation: bool = False """Stop after computing and saving attribution scores, before the leave-k-out retraining loop. Useful for score-only MAGIC runs.""" + + skip_metagradient: bool = False + """Skip the MAGIC backward pass entirely, using the trainer only. Attribution + scores are filled with zeros. Useful when another method (SOURCE, EK-FAC) + supplies the scores and MAGIC is wanted purely for its checkpoint schedule + and optimizer-state export.""" diff --git a/bergson/magic/metasmoothness.py b/bergson/magic/metasmoothness.py index e96adc78..1e20beba 100644 --- a/bergson/magic/metasmoothness.py +++ b/bergson/magic/metasmoothness.py @@ -46,6 +46,39 @@ def metasmoothness_score( return float((d / total * s1 * s2).sum()) +def metasmoothness_group_scores( + theta0: torch.Tensor, + theta_h: torch.Tensor, + theta_2h: torch.Tensor, + group_ids: torch.Tensor, + group_names: list[str], +) -> dict: + """Per-optimizer-path breakdown of the metasmoothness score. + + Muon routes 2D parameters through Newton-Schulz and 1D parameters through + an AdamW fallback; ``eps_root`` only enters the AdamW branch. Since the + global score weights coordinates by their L1 movement, the 2D group (99.9% + of GPT-2's parameters) can dominate it and mask a non-smooth 1D group. + + The decomposition is exact: ``score == sum(share_g * score_g)``. + """ + d = (theta_2h - theta0).abs() + total = d.sum() + contrib = d * torch.sign(theta_h - theta0) * torch.sign(theta_2h - theta_h) + + out = {} + for gid, name in enumerate(group_names): + m = group_ids == gid + dg = d[m].sum() + out[name] = { + "score": float(contrib[m].sum() / dg) if dg > 0 else 1.0, + "movement_share": float(dg / total) if total > 0 else 0.0, + "movement_l1": float(dg), + "numel": int(m.sum()), + } + return out + + def metasmoothness_worker( global_rank: int, rank: int, @@ -97,6 +130,9 @@ def metasmoothness_worker( v[-weight_pad_count:] = 0.0 thetas: list[torch.Tensor] = [] + theta_init: torch.Tensor | None = None + group_ids: torch.Tensor | None = None + group_names: list[str] = [] for k in range(3): weights = 1.0 + run_cfg.fd_step * k * v if pad_count: @@ -108,6 +144,16 @@ def metasmoothness_worker( trainer, fwd_state, model = prepare_trainer(run_cfg, rank, schedule) fwd_state.detach_() + if k == 0 and global_rank == 0: + theta_init = torch.cat( + [p.detach().float().cpu().flatten() for p in fwd_state.params.values()] + ) + if run_cfg.save_models: + os.makedirs(run_cfg.run_path, exist_ok=True) + torch.save( + {n: p.detach().float().cpu() for n, p in fwd_state.params.items()}, + os.path.join(run_cfg.run_path, "theta_init.pt"), + ) fwd_state = trainer.train(fwd_state, stream, inplace=True, fsdp=run_cfg.fsdp) if global_rank == 0: @@ -115,6 +161,27 @@ def metasmoothness_worker( [p.detach().float().cpu().flatten() for p in fwd_state.params.values()] ) thetas.append(theta) + if group_ids is None and run_cfg.optimizer == "muon": + # Muon's own split: ndim==2 -> Newton-Schulz, else AdamW fallback + # (the only branch eps_root enters). See optim.py muon(). + group_names = ["muon_2d", "adamw_1d"] + group_ids = torch.cat( + [ + torch.full( + (p.numel(),), 0 if p.ndim == 2 else 1, dtype=torch.uint8 + ) + for p in fwd_state.params.values() + ] + ) + if run_cfg.save_models and k == 0: + # Only the unperturbed run: the k>0 models are finite-difference + # probes, not models anyone wants to reuse. + path = os.path.join(run_cfg.run_path, "theta_final.pt") + torch.save( + {n: p.detach().float().cpu() for n, p in fwd_state.params.items()}, + path, + ) + print(f"[metasmoothness] saved {path}") print(f"[metasmoothness] finished training {k + 1}/3 (w = 1 + {k}*h*v)") del trainer, fwd_state, model @@ -127,7 +194,27 @@ def metasmoothness_worker( "direction_seed": run_cfg.direction_seed, "total_movement_l1": movement, } + if theta_init is not None: + # Relative update norms of the unperturbed run, matching the ΔL1/ΔL2 + # convention in LDS_RESULTS (‖θ_final − θ_init‖ / ‖θ_init‖). + delta = thetas[0] - theta_init + result["delta_l1"] = float(delta.abs().sum() / theta_init.abs().sum()) + result["delta_l2"] = float(delta.norm() / theta_init.norm()) + print( + f"[metasmoothness] deltaL1 = {result['delta_l1']:.4g} " + f"deltaL2 = {result['delta_l2']:.4g}" + ) print(f"[metasmoothness] score = {score:.4f} (h={run_cfg.fd_step})") + if group_ids is not None: + groups = metasmoothness_group_scores( + thetas[0], thetas[1], thetas[2], group_ids, group_names + ) + result["groups"] = groups + for name, g in groups.items(): + print( + f"[metasmoothness] {name}: score = {g['score']:.4f} " + f"share = {g['movement_share']:.4f} numel = {g['numel']:,}" + ) os.makedirs(run_cfg.run_path, exist_ok=True) with open(os.path.join(run_cfg.run_path, "metasmoothness.json"), "w") as f: json.dump(result, f, indent=2) diff --git a/bergson/magic/trainer.py b/bergson/magic/trainer.py index 6a92c33a..9acf7015 100644 --- a/bergson/magic/trainer.py +++ b/bergson/magic/trainer.py @@ -106,7 +106,9 @@ def result(self): return result -def next_save_index(current: int, n: int, save_mode: MagicSaveMode) -> int: +def next_save_index( + current: int, n: int, save_mode: MagicSaveMode, save_interval: int = 0 +) -> int: """Index of the checkpoint following the one saved at step `current`. `n` is the total number of training steps. The result is always strictly @@ -117,6 +119,10 @@ def next_save_index(current: int, n: int, save_mode: MagicSaveMode) -> int: case "all": # Save every step return current + 1 + case "interval": + if save_interval <= 0: + raise ValueError("save_mode='interval' requires save_interval > 0.") + return current + save_interval case "sqrt": chunk_size = math.isqrt(n) @@ -555,6 +561,7 @@ def train( max_grad_norm: float | None = None, grad_accum_steps: int = 1, optimizer_cfg: dict | None = None, + save_interval: int = 0, ) -> TrainerState: """Train the model on the given data stream, starting from the given state. @@ -583,6 +590,7 @@ def train( optimizer_cfg: When set (to the optimizer's betas/eps/eps_root), write each checkpoint's second moments to ``optimizer.pt``, tagged with the step and these hyperparameters. AdamW only. + save_interval: Snapshot spacing in steps for ``save_mode="interval"``. Returns: The final trainer state after training. @@ -634,7 +642,7 @@ def train( pending_save = state.save(p, debug_pbar=pbar if debug else None) - next_save = next_save_index(next_save, n, save_mode) + next_save = next_save_index(next_save, n, save_mode, save_interval) x = data[i] state = self.step( @@ -653,6 +661,23 @@ def train( if pending_save is not None: pending_save.result() + # Snapshots are written BEFORE each step, so the post-training state has + # no snapshot of its own; write it when the cadence lands exactly on n. + # Interval mode only: the other modes' backward replay indexes the data + # stream by checkpoint, so an extra trailing snapshot runs it past the end. + if save_dir and save_mode == "interval" and next_save == n: + p = os.path.join(save_dir, f"step_{n}.ckpt") + if optimizer_cfg is not None: + os.makedirs(p, exist_ok=True) + save_second_moments_as_optimizer_pt( + self.model, + state.opt_state, + os.path.join(p, "optimizer.pt"), + step=n, + **optimizer_cfg, + ) + state.save(p).result() + return state def save_backward_state(self, bwd_state, path, expected_idx, last_idx): diff --git a/bergson/validate.py b/bergson/validate.py index 447613d2..ba824707 100644 --- a/bergson/validate.py +++ b/bergson/validate.py @@ -80,6 +80,62 @@ def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]: return scores, False +def lds_from_precomputed_subsets( + score_path: str, + masks_path: str, + losses_path: str, + summary_path: str | None = None, +) -> float: + """LDS against a precomputed subset-retrain ground truth. + + ``masks_path``: ``[M, N_train]`` 0/1 subset-membership array (row ``j`` + marks the training rows the ``j``-th bank model was trained ON); + ``losses_path``: ``[M, Q]`` per-subset query losses (averaged over retrain + seeds). Follows Park et al. 2023 / Bae et al. 2024 Eq. 26-27: the group + prediction for subset ``j`` is the summed score of its members, and the LDS + is the per-query Spearman between predictions and subset losses, averaged + over queries. ``load_attribution_scores`` orients scores to the loss-diff + convention, under which the prediction correlates POSITIVELY with the + subset loss, so the returned mean Spearman is comparable to published LDS. + """ + scores, _ = load_attribution_scores(score_path) + if scores.ndim == 1: + scores = scores[:, None] + if scores.ndim != 2: + raise ValueError( + f"Expected per-document scores [N_train, Q]; got shape " + f"{tuple(scores.shape)}." + ) + masks = np.asarray(torch.load(masks_path, weights_only=False), dtype=np.float64) + losses = np.asarray(torch.load(losses_path, weights_only=False), dtype=np.float64) + if masks.shape[1] != scores.shape[0]: + raise ValueError( + f"masks cover {masks.shape[1]} training rows but scores have " + f"{scores.shape[0]}; the score run must use the identical dataset " + "ordering the masks index into." + ) + if losses.shape != (masks.shape[0], scores.shape[1]): + raise ValueError( + f"losses shape {losses.shape} does not match " + f"[num_subsets={masks.shape[0]}, num_queries={scores.shape[1]}]." + ) + + preds = masks @ scores.double().numpy() + per_query = np.array( + [spearmanr(preds[:, q], losses[:, q]).statistic for q in range(losses.shape[1])] + ) + mean_lds = float(np.nanmean(per_query)) + + if summary_path is not None: + writer = CSVWriter(summary_path, ["query", "spearman_corr"]) + for q, r in enumerate(per_query): + writer.writerow(q, float(r)) + writer.writerow("mean", mean_lds) + writer.close() + print(f"LDS (mean Spearman over {losses.shape[1]} queries): {mean_lds:.4f}") + return mean_lds + + def bank_loss_cache_key( run_cfg: ValidationConfig, multi_query: bool, num_subsets: int ) -> str: diff --git a/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm.yaml b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm.yaml new file mode 100644 index 00000000..d6303c17 --- /dev/null +++ b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm.yaml @@ -0,0 +1,83 @@ +# SOURCE (approximate unrolling, Bae et al. 2024) for the lotus GPT-2 base +# model against the 50 test queries, validated on the 400-model lotus bank. +# Plain SGD variant (no optimizer preconditioning). +# +# fisher_normalization: document -- the segment EK-FAC eigenvalues are a mean +# over the segment's (checkpoint, document) pairs, matching kronfluence's +# lambda_matrix / num_lambda_processed. This is the replication setting. +# +# Baseline for comparison: runs/lotus_source_q50_damp0 (unnormalized sums, +# LDS rho=0.387 +/- 0.039). + +# Run with `bergson examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm.yaml` + +steps: + - approxunrolling: + index_cfg: + run_path: runs/lotus_source_q50_docnorm + # Tokenizer + final-model reference; per-checkpoint passes load the + # model dirs below. + model: runs/lotus/retrained/base + # In the approxunrolling pipeline `overwrite` doubles as resume: + # steps whose output dirs exist are skipped after a crash. + overwrite: true + precision: fp32 + # GPT-2's max context length is 1024. + token_batch_size: 1024 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 # <-- adjust to this node's GPU count + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + # Doc-level scoring (one doc per row, no packing): packed 512-token + # chunks contaminate per-doc gradients across doc boundaries and + # were found to suppress LDS (see gpt2_lotus_trackstar50q_docs). + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + approx_unrolling_cfg: + checkpoints: + - runs/lotus_source_q50/models/step_51 + - runs/lotus_source_q50/models/step_102 + - runs/lotus_source_q50/models/step_136 + - runs/lotus_source_q50/models/step_187 + - runs/lotus_source_q50/models/step_238 + - runs/lotus_source_q50/models/step_288 + segments: 3 + fisher_normalization: document + lr_list: [4.857423e-04, 5.616667e-04, 2.416667e-04] + step_size_list: [97, 96, 96] + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + # One score column per query (50 total), like the EK-FAC and + # TrackStar runs, for multi-query LDS validation. + query_aggregation: none + + - validate: + run_path: runs/lotus_source_q50_docnorm_validate + model: gpt2 + overwrite: true + + scores: runs/lotus_source_q50_docnorm/scores + retrained_dir: runs/lotus + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 512 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + + batch_size: 64 diff --git a/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm_validate.yaml b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm_validate.yaml new file mode 100644 index 00000000..bda9d939 --- /dev/null +++ b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_docnorm_validate.yaml @@ -0,0 +1,18 @@ +# Validation only: the SOURCE scores for the per-document run already exist. + +steps: +- validate: + run_path: runs/lotus_source_q50_docnorm_validate + model: gpt2 + overwrite: true + scores: runs/lotus_source_q50_docnorm/scores + retrained_dir: runs/lotus + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: train + chunk_length: 512 + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: test[1:51] + chunk_length: 0 + batch_size: 64 diff --git a/examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml new file mode 100644 index 00000000..9568a898 --- /dev/null +++ b/examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml @@ -0,0 +1,83 @@ +# SOURCE (approximate unrolling, Bae et al. 2024) for the lotus GPT-2 base +# model against the 50 test queries, validated on the 400-model lotus bank. +# Plain SGD variant (no optimizer preconditioning). +# +# fisher_normalization: token -- as the document variant, but dividing by the +# token count instead. LambdaCollector sums over documents, so this is the +# per-document mean scaled by ~1/mean_tokens_per_doc. +# +# Baseline for comparison: runs/lotus_source_q50_damp0 (unnormalized sums, +# LDS rho=0.387 +/- 0.039). + +# Run with `bergson examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml` + +steps: + - approxunrolling: + index_cfg: + run_path: runs/lotus_source_q50_toknorm + # Tokenizer + final-model reference; per-checkpoint passes load the + # model dirs below. + model: runs/lotus/retrained/base + # In the approxunrolling pipeline `overwrite` doubles as resume: + # steps whose output dirs exist are skipped after a crash. + overwrite: true + precision: fp32 + # GPT-2's max context length is 1024. + token_batch_size: 1024 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 # <-- adjust to this node's GPU count + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + # Doc-level scoring (one doc per row, no packing): packed 512-token + # chunks contaminate per-doc gradients across doc boundaries and + # were found to suppress LDS (see gpt2_lotus_trackstar50q_docs). + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + approx_unrolling_cfg: + checkpoints: + - runs/lotus_source_q50/models/step_51 + - runs/lotus_source_q50/models/step_102 + - runs/lotus_source_q50/models/step_136 + - runs/lotus_source_q50/models/step_187 + - runs/lotus_source_q50/models/step_238 + - runs/lotus_source_q50/models/step_288 + segments: 3 + fisher_normalization: token + lr_list: [4.857423e-04, 5.616667e-04, 2.416667e-04] + step_size_list: [97, 96, 96] + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + # One score column per query (50 total), like the EK-FAC and + # TrackStar runs, for multi-query LDS validation. + query_aggregation: none + + - validate: + run_path: runs/lotus_source_q50_toknorm_validate + model: gpt2 + overwrite: true + + scores: runs/lotus_source_q50_toknorm/scores + retrained_dir: runs/lotus + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 512 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[1:51]" + chunk_length: 0 + + batch_size: 64 diff --git a/examples/replication/wikitext_gpt2_ekfac.yaml b/examples/replication/wikitext_gpt2_ekfac.yaml new file mode 100644 index 00000000..59de6aa1 --- /dev/null +++ b/examples/replication/wikitext_gpt2_ekfac.yaml @@ -0,0 +1,56 @@ +# EK-FAC influence functions at the final checkpoint of the WikiText-2 +# replication run (examples/replication/wikitext_gpt2_train.yaml), scored +# against all 481 validation queries. +# +# Reference: kronfluence examples/wikitext reports LDS 0.44 against its shipped +# ground truth (masks.pt 100x4656 alpha=0.5 keep-masks, losses.pt 100x481 +# five-seed-averaged subset losses). Running this pipeline on kronfluence's own +# released model gives 0.4431, so the EK-FAC implementation itself is validated; +# this config checks that a model trained by OUR trainer lands in the same place. +# +# damping_factor 1e-8 is effectively undamped, matching the reference +# implementation, which adds damping only via its heuristic when unset. +# +# Run with: +# PYTHONPATH=$PWD python -m bergson examples/replication/wikitext_gpt2_ekfac.yaml + +steps: + - ekfac: + index_cfg: + run_path: runs/wikitext_repro/if_scores + model: runs/wikitext_repro/train/retrained/base + overwrite: true + precision: fp32 + token_batch_size: 1024 + max_batch_size: 64 + # GPT-2 ties lm_head to wte; tracking both double-counts the embedding. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 + nnode: 1 + data: + dataset: runs/wikitext_repro/data + split: train + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + score_cfg: + batch_size: 1024 + query_batch_size: 64 + + preprocess_cfg: + unit_normalize: false + + hessian_pipeline_cfg: + query: + dataset: runs/wikitext_repro/data + split: validation + chunk_length: 0 + # One score column per validation query, for per-query LDS. + query_aggregation: none + inversion_cfg: + damping_factor: 1.0e-8 diff --git a/examples/replication/wikitext_gpt2_source.yaml b/examples/replication/wikitext_gpt2_source.yaml new file mode 100644 index 00000000..b118dd47 --- /dev/null +++ b/examples/replication/wikitext_gpt2_source.yaml @@ -0,0 +1,76 @@ +# SOURCE on the WikiText-2 replication run, following Bae et al. 2024. +# +# The paper's WikiText-2/GPT-2 experiment is AdamW only (App. B.1, p. 28), so +# this uses the preconditioned variant of App. C, with the Eq-43 hybrid from +# App. D: "use the diagonal Hessian approximation for computing the matrix +# exponential in Equation 39 and Equation 43. Note that we still use the EK-FAC +# factors to compute H^-1 g in Equation 43." +# +# L=3 segments over C=6 checkpoints (Figure 6 uses L=3), constant lr 3e-5, +# 582 steps per segment. Queries: all 481 validation chunks, one score column +# each, for per-query LDS against the same ground truth as the EK-FAC config. +# +# adam_hybrid_damping 1e-8 keeps the chained EK-FAC inverse effectively +# undamped, matching the reference implementation. +# +# fisher_normalization: document is required, not cosmetic. The eigenfunctions +# are not scale-invariant, and Eq. 1 defines the empirical risk as a mean over +# data points, so the segment eigenvalues must be a per-document expectation. +# Token units put the mask argument in its dead-linear range (LDS 0.125 vs +# 0.383 in an earlier run of this experiment). +# +# The checkpoint dirs come from exporting the trainer's DCP snapshots: +# python -c "from bergson.utils.trainer_export import export_checkpoints; \ +# export_checkpoints('runs/wikitext_repro/train', \ +# steps=[291,582,873,1164,1455,1746], overwrite=True)" +# +# Run with: +# PYTHONPATH=$PWD python -m bergson examples/replication/wikitext_gpt2_source.yaml + +steps: + - approxunrolling: + index_cfg: + run_path: runs/wikitext_repro/source_scores + model: runs/wikitext_repro/train/retrained/base + overwrite: true + precision: fp32 + token_batch_size: 1024 + max_batch_size: 64 + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 + nnode: 1 + data: + dataset: runs/wikitext_repro/data + split: train + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + approx_unrolling_cfg: + checkpoints: + - runs/wikitext_repro/train/exported/checkpoint-291 + - runs/wikitext_repro/train/exported/checkpoint-582 + - runs/wikitext_repro/train/exported/checkpoint-873 + - runs/wikitext_repro/train/exported/checkpoint-1164 + - runs/wikitext_repro/train/exported/checkpoint-1455 + - runs/wikitext_repro/train/exported/checkpoint-1746 + segments: 3 + lr_list: [3.0e-5, 3.0e-5, 3.0e-5] + step_size_list: [582, 582, 582] + query: + dataset: runs/wikitext_repro/data + split: validation + chunk_length: 0 + query_aggregation: none + query_batch_size: 64 + + # App. C preconditioned variant + the App. D Eq-43 hybrid. + use_adam_preconditioner: true + adam_segment_hybrid: true + adam_hybrid_damping: 1.0e-8 + + fisher_normalization: document diff --git a/examples/replication/wikitext_gpt2_train.yaml b/examples/replication/wikitext_gpt2_train.yaml new file mode 100644 index 00000000..9d6e9130 --- /dev/null +++ b/examples/replication/wikitext_gpt2_train.yaml @@ -0,0 +1,67 @@ +# Bae et al. 2024 WikiText-2 / GPT-2 training recipe (Appendix B.1, p. 28): +# "we fine-tuned the GPT-2 model using the WikiText-2 dataset. We followed the +# training script from the Transformer library but set the maximum sequence +# length to 512. During fine-tuning with AdamW, we saved 6 intermediate +# checkpoints for data attribution. The learning rate, weight decay, and +# batch size were set to 3e-5, 1e-2, and 8, respectively." +# +# Matches kronfluence examples/wikitext (3 epochs, seed 1004) so the resulting +# model is comparable to the one their reported LDS numbers were computed on. +# +# train_mode: true is load-bearing. kronfluence's train.py calls model.train(), +# so GPT-2's default 0.1 resid/embd/attn dropout is active during fine-tuning. +# Training with it off produces a measurably different model: EK-FAC LDS 0.372 +# on a dropout-free model vs 0.443 on kronfluence's own released model. +# +# 4656 train rows / 8 per batch = 582 steps per epoch, 1746 total; a checkpoint +# every 291 steps gives the paper's 6. +# +# Run with: +# PYTHONPATH=$PWD python -m bergson examples/replication/wikitext_gpt2_train.yaml + +steps: + - magic: + run_path: runs/wikitext_repro/train + model: gpt2 + overwrite: true + precision: fp32 + seed: 1004 + + # MAGIC is used purely as a trainer here; no metagradient is needed. + skip_metagradient: true + skip_validation: true + cleanup_ckpts: false + + data: + dataset: runs/wikitext_repro/data + split: train + chunk_length: 0 + query: + dataset: runs/wikitext_repro/data + split: validation + chunk_length: 0 + + distributed: + nproc_per_node: 1 + nnode: 1 + + optimizer: adamw + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-8 + eps_root: 0.0 + weight_decay: 0.01 + batch_size: 8 + num_epochs: 3 + lr_schedule: + lr_scheduler_type: constant + lr: 3.0e-5 + + # Dropout active, as in kronfluence's train.py. + train_mode: true + + save_mode: interval + save_interval: 291 + # SOURCE's Adam variant needs the optimizer state at every checkpoint. + save_optimizer_state: all + save_models: true diff --git a/experiments/source_fisher_normalization/chain_toknorm.sh b/experiments/source_fisher_normalization/chain_toknorm.sh new file mode 100755 index 00000000..d4875d7c --- /dev/null +++ b/experiments/source_fisher_normalization/chain_toknorm.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# Wait for the docnorm SOURCE run to finish, then start the toknorm one. +# Both need all 8 GPUs, so they cannot overlap. +# +# Bounded: gives up at DEADLINE_S, and exits early if the docnorm process +# disappears without producing scores. +set -u + +REPO=/mnt/ssd-1/lucia/bergson-damping +cd "$REPO" || exit 1 + +DEADLINE_S=$((6 * 3600)) +DOC_SCORES=runs/lotus_source_q50_docnorm/scores.npy +DOC_VALIDATE=runs/lotus_source_q50_docnorm_validate/summary.csv +TOK_CFG=examples/method_comparison/gpt2_wikitext_lotus_source_q50_toknorm.yaml + +start=$(date +%s) +while true; do + now=$(date +%s) + elapsed=$((now - start)) + if [ "$elapsed" -gt "$DEADLINE_S" ]; then + echo "CHAIN: deadline of ${DEADLINE_S}s reached; not launching toknorm" + exit 1 + fi + + if [ -f "$DOC_VALIDATE" ]; then + echo "CHAIN: docnorm validation complete after ${elapsed}s" + break + fi + + if ! pgrep -f "bergson.*docnorm" >/dev/null 2>&1; then + # Process gone. Only proceed if it actually got to the end. + if [ -f "$DOC_VALIDATE" ]; then + echo "CHAIN: docnorm finished after ${elapsed}s" + break + fi + echo "CHAIN: docnorm process gone after ${elapsed}s without $DOC_VALIDATE" \ + "(scores.npy exists: $([ -f "$DOC_SCORES" ] && echo yes || echo no)); aborting" + exit 1 + fi + + sleep 60 +done + +CMD="PYTHONPATH=$REPO python -m bergson $TOK_CFG" +echo "CHAIN: launching toknorm" +echo "CMD: $CMD" +echo "CMD: $CMD" > logs/source_toknorm.log +PYTHONPATH="$REPO" python -m bergson "$TOK_CFG" >> logs/source_toknorm.log 2>&1 +echo "CHAIN: toknorm exited with $?" diff --git a/experiments/source_fisher_normalization/compare_lds.py b/experiments/source_fisher_normalization/compare_lds.py new file mode 100644 index 00000000..17a812be --- /dev/null +++ b/experiments/source_fisher_normalization/compare_lds.py @@ -0,0 +1,81 @@ +"""Compare LDS across SOURCE Fisher-normalization settings. + +Reads each validate run's summary.csv (one row per query) and reports the mean +correlation with a 95% CI over queries, which is the convention behind the +published `0.387 +/- 0.039`. + + python experiments/source_fisher_normalization/compare_lds.py \ + none=runs/lotus_source_q50_damp0_validate \ + document=runs/lotus_source_q50_docnorm_validate \ + token=runs/lotus_source_q50_toknorm_validate +""" + +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy.stats import wilcoxon + +CI95 = 1.959963985 + + +def summarize(path: Path) -> dict: + df = pd.read_csv(path / "summary.csv") + out: dict = {"n_queries": len(df)} + for col, key in (("spearman_corr", "rho"), ("pearson_corr", "r")): + v = df[col].to_numpy() + out[key] = v.mean() + out[f"{key}_ci"] = CI95 * v.std(ddof=1) / np.sqrt(len(v)) + out[f"{key}_raw"] = v + return out + + +def main() -> None: + args = sys.argv[1:] + if not args: + raise SystemExit(__doc__) + + runs = {} + for arg in args: + label, _, path = arg.partition("=") + p = Path(path) / "summary.csv" + if not p.exists(): + print(f"skip {label}: no {p}") + continue + runs[label] = summarize(Path(path)) + + print(f"\n{'normalization':<14}{'rho':>18}{'r':>18}{'queries':>9}") + print("-" * 59) + for label, s in runs.items(): + print( + f"{label:<14}" + f"{s['rho']:>10.3f} +/- {s['rho_ci']:.3f}" + f"{s['r']:>10.3f} +/- {s['r_ci']:.3f}" + f"{s['n_queries']:>9}" + ) + + # Paired over the same 50 queries, so a paired test is the informative one. + if "none" in runs: + base = runs["none"] + for label, s in runs.items(): + if label == "none" or len(s["rho_raw"]) != len(base["rho_raw"]): + continue + delta = s["rho_raw"] - base["rho_raw"] + stat = wilcoxon(delta) + print( + f"\n{label} vs none: mean drho={delta.mean():+.3f} " + f"(won {int((delta > 0).sum())}/{len(delta)} queries, " + f"wilcoxon p={stat.pvalue:.2g})" + ) + + print("\nLaTeX rows:") + for label, s in runs.items(): + print( + f"SOURCE ({label}) & ${s['rho']:.3f} \\pm {s['rho_ci']:.3f}$ " + f"& ${s['r']:.3f} \\pm {s['r_ci']:.3f}$ \\\\" + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/source_fisher_normalization/setup_run.py b/experiments/source_fisher_normalization/setup_run.py new file mode 100644 index 00000000..62efa7f5 --- /dev/null +++ b/experiments/source_fisher_normalization/setup_run.py @@ -0,0 +1,107 @@ +"""Seed a SOURCE run dir from an existing one, reusing everything the Fisher +normalization does not change. + +Only the segment lambdas (step 4) and everything downstream depend on +``fisher_normalization``; the per-checkpoint covariances/lambdas, the segment +eigenvectors and the query gradients do not. Those are hardlinked from the +baseline run so a normalization sweep costs no extra Hessian passes. + +Derived outputs are deliberately NOT linked: the pipeline truncates them in +place, which would corrupt the baseline through a shared inode. + + python experiments/source_fisher_normalization/setup_run.py \ + runs/lotus_source_q50_damp0 runs/lotus_source_q50_docnorm +""" + +import argparse +import os +from pathlib import Path + +# Per-segment inputs that are independent of the lambda normalization. +SEGMENT_REUSE = ( + "activation_sharded", + "gradient_sharded", + "eigen_activation_sharded", + "eigen_gradient_sharded", + "total_processed.pt", +) +CKPT_REUSE = ( + "activation_sharded", + "gradient_sharded", + "averaged_ev_correct_sharded", + "total_processed.pt", +) + + +def link_tree(src: Path, dst: Path) -> int: + """Hardlink every file under ``src`` into ``dst``. Returns the file count.""" + n = 0 + if src.is_file(): + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + os.link(src, dst) + return 1 + for path in src.rglob("*"): + if path.is_dir(): + continue + out = dst / path.relative_to(src) + out.parent.mkdir(parents=True, exist_ok=True) + if not out.exists(): + os.link(path, out) + n += 1 + return n + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("baseline", type=Path) + ap.add_argument("target", type=Path) + args = ap.parse_args() + + baseline, target = args.baseline, args.target + if not baseline.exists(): + raise SystemExit(f"baseline run {baseline} not found") + if target.exists(): + raise SystemExit(f"{target} already exists; remove it first") + + total = 0 + total += link_tree(baseline / "query", target / "query") + + for seg_dir in sorted(baseline.glob("segment_*")): + seg = seg_dir.name + # A Hessian method dir, as opposed to a derived dir (scores, + # query_grad_*) that also happens to hold a total_processed.pt. + method_dirs = [ + d for d in seg_dir.iterdir() if (d / "eigen_activation_sharded").is_dir() + ] + for method_dir in method_dirs: + for name in SEGMENT_REUSE: + src = method_dir / name + if src.exists(): + total += link_tree(src, target / seg / method_dir.name / name) + + for ckpt_dir in sorted(seg_dir.glob("ckpt_*")): + for method_dir in ckpt_dir.iterdir(): + if not method_dir.is_dir(): + continue + # kfac.part holds the token count the backfill reads. + names = ( + CKPT_REUSE + if not method_dir.name.endswith(".part") + else ("total_processed.pt",) + ) + for name in names: + src = method_dir / name + if src.exists(): + total += link_tree( + src, target / seg / ckpt_dir.name / method_dir.name / name + ) + + print(f"hardlinked {total} files from {baseline} -> {target}") + print( + "not linked (regenerated): eigenvalue_correction_sharded, query_grad_*, scores" + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/wikitext_replication/run_chain.py b/experiments/wikitext_replication/run_chain.py new file mode 100644 index 00000000..ec87a8ea --- /dev/null +++ b/experiments/wikitext_replication/run_chain.py @@ -0,0 +1,64 @@ +"""Export checkpoints -> EK-FAC IF -> SOURCE -> LDS for the WikiText-2 replication. + +Assumes examples/replication/wikitext_gpt2_train.yaml has already run. +Ground truth is kronfluence's shipped masks/losses (see GROUND_TRUTH below). + + PYTHONPATH=$PWD python experiments/wikitext_replication/run_chain.py + +Last run (2026-07-31, 8x GPU): EK-FAC IF 0.4451 (kronfluence reports 0.44), +SOURCE 0.4506. SOURCE > IF reproduces the ordering in the paper's Figure 6. +""" + +import subprocess +import sys +from pathlib import Path + +from bergson.utils.trainer_export import export_checkpoints +from bergson.validate import lds_from_precomputed_subsets + +RUN = Path("runs/wikitext_repro") +GROUND_TRUTH = Path("/mnt/ssd-1/lucia/bergson-source-repro/runs/wikitext_kron") +STEPS = [291, 582, 873, 1164, 1455, 1746] + + +def run(cmd: str) -> None: + print(f"\n=== RUN: {cmd}", flush=True) + subprocess.run(cmd, shell=True, check=True) + + +def main() -> None: + masks, losses = GROUND_TRUTH / "masks.pt", GROUND_TRUTH / "losses.pt" + for p in (masks, losses): + if not p.exists(): + sys.exit(f"missing ground truth: {p}") + + exported = export_checkpoints(RUN / "train", steps=STEPS, overwrite=True) + print(f"exported: {[str(p) for p in exported]}", flush=True) + + run( + "python -m bergson examples/replication/wikitext_gpt2_ekfac.yaml" + f" > {RUN}/if.log 2>&1" + ) + print("EK-FAC IF: ", end="") + lds_from_precomputed_subsets( + str(RUN / "if_scores/scores"), + str(masks), + str(losses), + summary_path=str(RUN / "if_lds.csv"), + ) + + run( + "python -m bergson examples/replication/wikitext_gpt2_source.yaml" + f" > {RUN}/source.log 2>&1" + ) + print("SOURCE: ", end="") + lds_from_precomputed_subsets( + str(RUN / "source_scores/scores"), + str(masks), + str(losses), + summary_path=str(RUN / "source_lds.csv"), + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_lds_precomputed_subsets.py b/tests/test_lds_precomputed_subsets.py new file mode 100644 index 00000000..e24a9de1 --- /dev/null +++ b/tests/test_lds_precomputed_subsets.py @@ -0,0 +1,78 @@ +"""LDS against a precomputed subset-retrain ground truth. + +``lds_from_precomputed_subsets`` is the scoring side of a published bank +(kronfluence's ``masks.pt``/``losses.pt``, say): the prediction for a subset is +the summed score of its members, and the LDS is the per-query Spearman against +the measured subset losses. These cover the group-prediction arithmetic, the +per-query averaging, and the shape guards that catch a scores/masks mismatch. +""" + +import numpy as np +import pytest +import torch + +from bergson.validate import lds_from_precomputed_subsets + + +@pytest.fixture +def bank(tmp_path): + """A 6-subset bank over 5 training rows and 2 queries, with losses made + exactly linear in the summed scores so the LDS is 1.0 by construction.""" + rng = np.random.default_rng(0) + masks = rng.integers(0, 2, size=(6, 5)).astype(np.float64) + scores = rng.normal(size=(5, 2)) + + torch.save(torch.from_numpy(masks), tmp_path / "masks.pt") + torch.save(torch.from_numpy(masks @ scores), tmp_path / "losses.pt") + np.save(tmp_path / "scores.npy", scores.astype(np.float32)) + return tmp_path + + +def test_perfectly_predictive_scores_score_one(bank): + lds = lds_from_precomputed_subsets( + str(bank / "scores.npy"), str(bank / "masks.pt"), str(bank / "losses.pt") + ) + assert lds == pytest.approx(1.0) + + +def test_sign_flipped_scores_score_minus_one(bank): + scores = -np.load(bank / "scores.npy") + np.save(bank / "flipped.npy", scores) + + lds = lds_from_precomputed_subsets( + str(bank / "flipped.npy"), str(bank / "masks.pt"), str(bank / "losses.pt") + ) + assert lds == pytest.approx(-1.0) + + +def test_summary_csv_has_one_row_per_query_plus_mean(bank): + out = bank / "lds.csv" + lds = lds_from_precomputed_subsets( + str(bank / "scores.npy"), + str(bank / "masks.pt"), + str(bank / "losses.pt"), + summary_path=str(out), + ) + + rows = out.read_text().strip().splitlines() + assert len(rows) == 1 + 2 + 1 # header, two queries, mean + assert rows[-1].startswith("mean,") + assert float(rows[-1].split(",")[1]) == pytest.approx(lds) + + +def test_mismatched_training_row_count_is_rejected(bank): + np.save(bank / "short.npy", np.zeros((4, 2), dtype=np.float32)) + + with pytest.raises(ValueError, match="masks cover 5 training rows"): + lds_from_precomputed_subsets( + str(bank / "short.npy"), str(bank / "masks.pt"), str(bank / "losses.pt") + ) + + +def test_mismatched_query_count_is_rejected(bank): + np.save(bank / "wide.npy", np.zeros((5, 3), dtype=np.float32)) + + with pytest.raises(ValueError, match="losses shape"): + lds_from_precomputed_subsets( + str(bank / "wide.npy"), str(bank / "masks.pt"), str(bank / "losses.pt") + ) diff --git a/tests/test_source_fisher_normalization.py b/tests/test_source_fisher_normalization.py new file mode 100644 index 00000000..559f0771 --- /dev/null +++ b/tests/test_source_fisher_normalization.py @@ -0,0 +1,83 @@ +"""Normalization of the segment EK-FAC eigenvalues fed to SOURCE's eigenfunctions. + +kronfluence divides its lambda matrix by a per-sample count before use +(``kronfluence/factor/config.py`` ``lambda_matrix.div_(NUM_LAMBDA_PROCESSED)``, +counted in ``module/tracker/factor.py`` as ``per_sample_gradient.size(0)``). +``LambdaCollector`` accumulates an unnormalized sum, so the segment aggregation +divides by the pooled count to match. +""" + +import json + +import pytest +import torch +from safetensors.torch import load_file, save_file + +from bergson.approx_unrolling.segment_aggregation import ( + LAMBDA_COUNTS_FILENAME, + _sum_my_shard, + lambda_denominator, + write_lambda_counts, +) + +DOCS, TOKENS = 4599, 2354682 + + +def _ckpt_dirs(tmp_path, n=2): + dirs = [] + for i in range(n): + d = tmp_path / f"ckpt_{i}" + d.mkdir() + write_lambda_counts(d, documents=DOCS, tokens=TOKENS) + dirs.append(d) + return dirs + + +def test_counts_roundtrip(tmp_path): + write_lambda_counts(tmp_path, documents=DOCS, tokens=TOKENS) + with open(tmp_path / LAMBDA_COUNTS_FILENAME) as f: + assert json.load(f) == {"documents": DOCS, "tokens": TOKENS} + + +@pytest.mark.parametrize( + "normalization, expected", + [("document", 2 * DOCS), ("token", 2 * TOKENS), ("none", 1.0)], +) +def test_denominator_pools_over_checkpoints(tmp_path, normalization, expected): + """The lambdas are summed over checkpoints, so the denominator sums too.""" + assert lambda_denominator(_ckpt_dirs(tmp_path), normalization) == expected + + +def test_denominator_reports_missing_counts(tmp_path): + d = tmp_path / "ckpt_0" + d.mkdir() + with pytest.raises(FileNotFoundError, match="fisher_normalization"): + lambda_denominator([d], "document") + + +def test_denominator_ignores_missing_counts_when_unnormalized(tmp_path): + """ "none" reproduces pre-normalization runs, which have no counts.json.""" + d = tmp_path / "ckpt_0" + d.mkdir() + assert lambda_denominator([d], "none") == 1.0 + + +def test_sum_my_shard_divides(tmp_path): + ins = [] + for i in range(2): + p = tmp_path / f"in_{i}.safetensors" + save_file({"w": torch.full((2, 3), float(i + 1))}, p) + ins.append(p) + + out = tmp_path / "out.safetensors" + _sum_my_shard(ins, out, device="cpu", divisor=6.0) + # (1 + 2) / 6 + assert torch.allclose(load_file(out)["w"], torch.full((2, 3), 0.5)) + + +def test_sum_my_shard_defaults_to_plain_sum(tmp_path): + p = tmp_path / "in.safetensors" + save_file({"w": torch.ones(2, 3)}, p) + out = tmp_path / "out.safetensors" + _sum_my_shard([p], out, device="cpu") + assert torch.allclose(load_file(out)["w"], torch.ones(2, 3)) diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index f562fef2..86292c69 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -22,13 +22,15 @@ write_lr_history, ) from bergson.config.config import ApproxUnrollingConfig, TrainingConfig +from bergson.config.config_io import save_run_config def _run_dir(tmp_path, **overrides): - """A bergson run directory with a config.yaml, as save_run_config writes it.""" - cfg = TrainingConfig(run_path=str(tmp_path), **overrides) - payload = {"magic": cfg.to_dict()} - (tmp_path / "config.yaml").write_text(yaml.safe_dump([payload])) + """A bergson run directory with a config.yaml. + + Written by ``save_run_config`` itself so the loader is pinned against the + real on-disk shape rather than a hand-rolled approximation of it.""" + save_run_config(TrainingConfig(run_path=str(tmp_path), **overrides), tmp_path) return tmp_path @@ -137,6 +139,23 @@ def test_missing_config_yaml_is_explained(tmp_path): load_training_config(tmp_path) +def test_load_reads_the_saved_steps_document(tmp_path): + """save_run_config wraps the step list in {steps, metadata}.""" + _run_dir(tmp_path, optimizer="adamw", adam_beta2=0.98) + + doc = yaml.safe_load((tmp_path / "config.yaml").read_text()) + assert set(doc) == {"steps", "metadata"} + assert load_training_config(tmp_path).adam_beta2 == pytest.approx(0.98) + + +def test_load_still_reads_a_bare_step_list(tmp_path): + """Runs saved before the {steps, metadata} wrapper stay readable.""" + cfg = TrainingConfig(run_path=str(tmp_path), optimizer="sgd", adam_beta1=0.8) + (tmp_path / "config.yaml").write_text(yaml.safe_dump([{"magic": cfg.to_dict()}])) + + assert load_training_config(tmp_path).adam_beta1 == pytest.approx(0.8) + + # ── LR history ────────────────────────────────────────────────────────────── @@ -504,3 +523,20 @@ def test_export_checkpoints_end_to_end(tmp_path): out = resolve(ApproxUnrollingConfig(checkpoints=[str(p) for p in exported])) assert out.model_path == model_name assert out.momentum == 0.0 # adamw + + +def test_resolve_ignores_a_non_training_config(tmp_path): + """A checkpoint dir's sibling config.yaml may belong to an attribution run. + + ``infer_trainer_run`` only checks that a config.yaml exists, so ``resolve`` + has to tolerate one it cannot read as a TrainingConfig. + """ + run = tmp_path / "attribution_run" + (run / "models" / "step_1").mkdir(parents=True) + (run / "config.yaml").write_text("steps:\n- approxunrolling:\n index_cfg: {}\n") + + cfg = ApproxUnrollingConfig(checkpoints=[str(run / "models" / "step_1")]) + resolved = resolve(cfg) + + assert resolved.momentum == 0.0 + assert resolved.model_path is None