From 3084a48c0a475720ac2ffae81581915d0d276031 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 12:20:07 +0000 Subject: [PATCH 01/13] feat: run SOURCE directly off a bergson training run SOURCE was built for HF Trainer checkpoints. It loads each checkpoint with from_pretrained, reads the step out of a `checkpoint-` directory name, and scrapes per-step learning rates from log_history.json / trainer_state.json. A bergson run satisfies none of that: it writes DCP `step_.ckpt` directories, and no HF artifacts at all. Every input had to be reconstructed by hand. Worse, one of those hand-typed inputs is silently wrong by default. bergson's SGD takes its momentum from TrainingConfig.adam_beta1 (see prepare_trainer), which defaults to 0.95, while the unrolling assumed no momentum. The terminal-velocity factor 1/(1-beta) (Bae et al. 2024, App. D.2) is therefore 20x, so every segment's lr*steps was understated 20-fold unless the user knew to look. `momentum` is now a config field derived from the run that trained the checkpoints. The pieces: - trainer_run.py resolves what a config leaves unset -- checkpoints, model_path, momentum -- from a bergson run's config.yaml. Called once at the top of the pipeline. - The trainer records its realized per-step LRs as log_history.json beside the checkpoints, in HF's shape. That is the file the LR math already checks first, so bergson runs join the existing path instead of branching it, and the schedule is exact rather than reconstructed. - trainer_export.py turns DCP checkpoints into `checkpoint-/` model dirs from_pretrained can load, carrying each step's optimizer state and the LR history across. Offline, so a run only pays the disk for trajectories being attributed. - _checkpoint_step also accepts `step_[.ckpt]`. Derivation is strictly a fallback: it only ever supplies a value the config did not give. Explicit lr_list/step_size_list/momentum/model_path/checkpoints always win, and with no `trainer_run` set, resolve() is a no-op -- so checkpoints from HF Trainer or any other trainer behave exactly as before. That is pinned by tests for each field, including that an explicit momentum of 0.0 survives (it is a real value, not "unset"). optimizer_placement.py lands here rather than in the ADAM SOURCE branch, since the export needs it and it is not Adam-specific. --- .../approx_unrolling/approx_unrolling_math.py | 30 +- bergson/approx_unrolling/pipeline.py | 6 + bergson/approx_unrolling/trainer_run.py | 208 +++++++++++++ bergson/config/config.py | 19 ++ bergson/magic/cli.py | 7 + bergson/utils/optimizer_placement.py | 204 +++++++++++++ bergson/utils/trainer_export.py | 134 ++++++++ tests/test_optimizer_placement.py | 208 +++++++++++++ tests/test_source_trainer_integration.py | 289 ++++++++++++++++++ 9 files changed, 1100 insertions(+), 5 deletions(-) create mode 100644 bergson/approx_unrolling/trainer_run.py create mode 100644 bergson/utils/optimizer_placement.py create mode 100644 bergson/utils/trainer_export.py create mode 100644 tests/test_optimizer_placement.py create mode 100644 tests/test_source_trainer_integration.py diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index eab33bc7..21ace167 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -29,19 +29,24 @@ def _checkpoint_step(p: str) -> int: """Extract the training step index for a checkpoint. Accepts a path whose final component is a ``checkpoint-`` directory - (what HF Trainer writes natively) or a bare ```` step directory. + (what HF Trainer writes natively), a ``step_[.ckpt]`` directory (what + the bergson trainer writes), or a bare ```` step directory. """ # TODO: Inferring the step from a checkpoint name via regex is brittle. # For now, we raise an error name = Path(p).name m = re.match(r"checkpoint-(\d+)$", name) + if m: + return int(m.group(1)) + m = re.match(r"step_(\d+)(?:\.ckpt)?$", name) if m: return int(m.group(1)) if name.isdigit(): return int(name) raise ValueError( f"Cannot infer a training step from checkpoint {p!r}: expected a path " - "ending in 'checkpoint-' or a bare '' step directory. Set both " + "ending in 'checkpoint-', 'step_[.ckpt]', or a bare '' step " + "directory. Set both " "`lr_list` and `step_size_list` on ApproxUnrollingConfig to specify " "the per-segment learning rate and step counts explicitly instead of " "inferring them from checkpoint names." @@ -52,11 +57,25 @@ def compute_lr_times_steps_per_segment( approx_unrolling_cfg: ApproxUnrollingConfig, ) -> list[float]: """Per-segment lr * K. Use ``lr_list * step_size_list`` if set on config; - else equal-partition log_history.json into segments and sum per-step LRs.""" + else equal-partition log_history.json into segments and sum per-step LRs. + + With SGD heavy-ball ``momentum`` beta, scale by the terminal velocity + 1/(1-beta) (Bae et al. 2024, App. D.2). ``momentum`` is taken as given when + set; ``None`` means it is derived from ``trainer_run`` (or 0.0 without one), + so checkpoints from another trainer are unaffected by the derivation. + """ cfg = approx_unrolling_cfg L = cfg.segments + + momentum = cfg.momentum if cfg.momentum is not None else 0.0 + if not 0.0 <= momentum < 1.0: + raise ValueError(f"momentum must be in [0, 1), got {momentum}.") + momentum_scale = 1.0 / (1.0 - momentum) + if cfg.lr_list and cfg.step_size_list: - return [lr * k for lr, k in zip(cfg.lr_list, cfg.step_size_list)] + return [ + lr * k * momentum_scale for lr, k in zip(cfg.lr_list, cfg.step_size_list) + ] per_segment = len(cfg.checkpoints) // L ckpt_steps = [_checkpoint_step(p) for p in cfg.checkpoints] @@ -75,7 +94,8 @@ def compute_lr_times_steps_per_segment( log_history = json.load(f)["log_history"] step_to_lr = {e["step"]: e["learning_rate"] for e in log_history} return [ - sum( + momentum_scale + * sum( step_to_lr.get(s, 0.0) for s in range(boundaries[l] + 1, boundaries[l + 1] + 1) ) diff --git a/bergson/approx_unrolling/pipeline.py b/bergson/approx_unrolling/pipeline.py index d5db4916..7a179cb8 100644 --- a/bergson/approx_unrolling/pipeline.py +++ b/bergson/approx_unrolling/pipeline.py @@ -52,6 +52,7 @@ aggregate_segment_covariances, aggregate_segment_lambdas, ) +from .trainer_run import resolve as resolve_trainer_run # Total number of steps in the full approximate unrolling pipeline. Used only for the # user-facing "Step k/N_TOTAL_STEPS:" prefix. Bump as steps land. @@ -82,6 +83,11 @@ def approx_unrolling_pipeline( """ logger = get_logger("approx_unrolling_pipeline") + # Fill in anything the config left unset from the bergson run it names. + # A no-op without `trainer_run`, and it never overrides an explicit value, + # so configs for checkpoints from other trainers are unaffected. + approx_unrolling_cfg = resolve_trainer_run(approx_unrolling_cfg) + n_ckpts = len(approx_unrolling_cfg.checkpoints) n_segments = approx_unrolling_cfg.segments if n_ckpts == 0: diff --git a/bergson/approx_unrolling/trainer_run.py b/bergson/approx_unrolling/trainer_run.py new file mode 100644 index 00000000..740874b1 --- /dev/null +++ b/bergson/approx_unrolling/trainer_run.py @@ -0,0 +1,208 @@ +"""Fill in SOURCE hyperparameters from a bergson training run. + +SOURCE was built around HF Trainer runs: it infers per-step learning rates from +``log_history.json`` / ``trainer_state.json`` and reads steps out of +``checkpoint-`` directory names. A bergson run has the same information in +its own form -- ``config.yaml`` plus the checkpoint schedule -- so this module +maps one onto the other. + +Everything here is a *fallback*. An explicitly configured field always wins, so +checkpoints produced by any other trainer keep working exactly as before by +setting ``lr_list``/``step_size_list``/``momentum`` by hand and leaving +``trainer_run`` empty. +""" + +import json +import os +from pathlib import Path +from typing import Any, Callable + +import yaml + +from ..config.config import ApproxUnrollingConfig, TrainingConfig +from ..config.config_io import CONFIG_FILENAME +from ..utils.logger import get_logger + +LR_HISTORY_FILENAME = "log_history.json" +"""Per-step learning rates, in HF Trainer's ``log_history`` shape. + +Written next to a bergson run's checkpoints so +:func:`~bergson.approx_unrolling.approx_unrolling_math.compute_lr_times_steps_per_segment` +picks it up through the path it already checks first, with no bergson-specific +branch in the LR math. +""" + +logger = get_logger(__name__) + + +def write_lr_history( + save_dir: str | Path, schedule: Callable[[int], float], num_steps: int +) -> Path: + """Record the run's realized per-step learning rates beside its checkpoints. + + ``schedule`` is the step -> lr callable the optimizer was built with, so + this is the exact schedule that ran rather than a reconstruction of it. + """ + save_dir = Path(save_dir) + save_dir.mkdir(parents=True, exist_ok=True) + history = [ + {"step": step, "learning_rate": float(schedule(step))} + for step in range(num_steps) + ] + path = save_dir / LR_HISTORY_FILENAME + with open(path, "w") as f: + json.dump(history, f) + return path + + +def load_training_config(trainer_run: str | Path) -> TrainingConfig: + """Load the ``TrainingConfig`` a bergson run was launched with. + + ``save_run_config`` writes ``config.yaml`` as ``{command_name: {...}}``, so + the single value is the run's config regardless of which command wrote it. + """ + path = Path(trainer_run) / CONFIG_FILENAME + if not path.is_file(): + raise FileNotFoundError( + f"{path} not found; trainer_run must point at a bergson run " + "directory (the one containing config.yaml and checkpoints/)." + ) + + with open(path) as f: + loaded = yaml.safe_load(f) + + # One-step configs are a list of {command: payload}; take the first payload. + if isinstance(loaded, list): + if not loaded: + raise ValueError(f"{path} is empty") + loaded = loaded[0] + if not isinstance(loaded, dict) or not loaded: + raise ValueError(f"{path} is not a bergson run config") + + payload: Any = next(iter(loaded.values())) + if not isinstance(payload, dict): + raise ValueError(f"{path} is not a bergson run config") + + return TrainingConfig.from_dict(payload, drop_extra_fields=True) + + +def derive_momentum(training_cfg: TrainingConfig) -> float: + """Heavy-ball momentum beta that a bergson run actually trained with. + + bergson's SGD passes ``adam_beta1`` as ``torchopt.sgd``'s momentum (see + ``prepare_trainer``), so the SOURCE-relevant beta is ``adam_beta1`` for SGD + and ``0.0`` for AdamW, whose own preconditioner already accounts for its + first moment. + """ + match training_cfg.optimizer: + case "sgd": + return float(training_cfg.adam_beta1) + case "adamw": + return 0.0 + case other: + # Muon routes 2D params through Newton-Schulz, which the unrolling + # derivation does not cover; don't invent a scaling for it. + logger.warning( + "Cannot derive a SOURCE momentum for optimizer %r; using 0.0. " + "Set ApproxUnrollingConfig.momentum explicitly if that is wrong.", + other, + ) + return 0.0 + + +def discover_checkpoints(trainer_run: str | Path) -> list[str]: + """The run's saved checkpoints, in training order. + + Prefers exported HF-format ``checkpoint-`` directories (what SOURCE can + actually load) and falls back to the trainer's native ``step_.ckpt``, so + the error surfaced to the caller names the export step rather than an empty + list. + """ + root = Path(trainer_run) + exported = sorted( + (p for p in root.glob("checkpoint-*") if p.is_dir()), + key=lambda p: int(p.name.split("-")[1]), + ) + if exported: + return [str(p) for p in exported] + + ckpt_dir = root / "checkpoints" + native = sorted( + (p for p in ckpt_dir.glob("step_*.ckpt") if p.is_dir()), + key=lambda p: int(p.name.removesuffix(".ckpt").removeprefix("step_")), + ) + if native: + raise FileNotFoundError( + f"{trainer_run} has {len(native)} native checkpoints under " + f"{ckpt_dir} but no exported checkpoint- directories. SOURCE " + "loads models with from_pretrained, so export them first with " + "bergson.utils.trainer_export.export_checkpoints(run_path, out_dir)." + ) + + raise FileNotFoundError( + f"No checkpoints found under {trainer_run}. Train with a save_mode " + "that writes them, then export with " + "bergson.utils.trainer_export.export_checkpoints." + ) + + +def resolve(cfg: ApproxUnrollingConfig) -> ApproxUnrollingConfig: + """Return ``cfg`` with unset fields filled in from ``cfg.trainer_run``. + + A no-op when ``trainer_run`` is empty, so configs for other trainers pass + through untouched. Fields already set are never overwritten -- this only + ever supplies a value the caller did not give. + """ + if not cfg.trainer_run: + # Nothing to derive from; only normalize the momentum sentinel. + if cfg.momentum is None: + cfg.momentum = 0.0 + return cfg + + training_cfg = load_training_config(cfg.trainer_run) + filled: list[str] = [] + + if not cfg.checkpoints: + cfg.checkpoints = discover_checkpoints(cfg.trainer_run) + filled.append(f"checkpoints ({len(cfg.checkpoints)})") + + if cfg.model_path is None: + cfg.model_path = training_cfg.model + filled.append(f"model_path={cfg.model_path!r}") + + if cfg.momentum is None: + cfg.momentum = derive_momentum(training_cfg) + filled.append(f"momentum={cfg.momentum}") + + if filled: + logger.info( + "Filled from trainer_run %s: %s", cfg.trainer_run, ", ".join(filled) + ) + + return cfg + + +def lr_history_path(cfg: ApproxUnrollingConfig) -> Path | None: + """Where a bergson run's LR history lives, if this config names one.""" + if not cfg.trainer_run: + return None + for candidate in ( + Path(cfg.trainer_run) / LR_HISTORY_FILENAME, + Path(cfg.trainer_run) / "checkpoints" / LR_HISTORY_FILENAME, + ): + if candidate.is_file(): + return candidate + return None + + +def _checkpoint_dir_step(path: str | os.PathLike) -> int | None: + """Step index named by a checkpoint dir, across the layouts we emit/accept.""" + name = Path(path).name + if name.startswith("checkpoint-") and name.removeprefix("checkpoint-").isdigit(): + return int(name.removeprefix("checkpoint-")) + stem = name.removesuffix(".ckpt") + if stem.startswith("step_") and stem.removeprefix("step_").isdigit(): + return int(stem.removeprefix("step_")) + if name.isdigit(): + return int(name) + return None diff --git a/bergson/config/config.py b/bergson/config/config.py index c86d004d..e21de0c0 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -786,6 +786,25 @@ class ApproxUnrollingConfig(Serializable): """Per-segment optimizer step count; length must == segments. If empty, inferred from per-checkpoint step counts.""" + trainer_run: str = "" + """Directory of a bergson training run, used to fill in what this config + does not state explicitly: ``checkpoints``, ``model_path`` and + ``momentum``. Leave empty for checkpoints from another trainer (e.g. HF + Trainer) and set those fields by hand -- anything set explicitly always + wins over what is read from the run.""" + + momentum: float | None = None + """SGD heavy-ball momentum beta used during training; scales each segment's + lr*steps by the terminal velocity 1/(1-beta) (Bae et al. 2024, App. D.2). + + ``None`` means "not set": derive it from ``trainer_run`` if one is given, + else fall back to ``0.0``. Set it explicitly (``0.0`` included) when + training happened outside bergson. + + Deriving matters because bergson's SGD takes its momentum from + ``TrainingConfig.adam_beta1``, which defaults to ``0.95`` -- assuming 0.0 + there understates lr*steps by 20x.""" + query: DataConfig = field(default_factory=DataConfig) """Query dataset spec; gradients computed at the final checkpoint.""" diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index 79970b05..319fd5ab 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -26,6 +26,7 @@ set_verbosity_error as hf_set_verbosity_error, ) +from ..approx_unrolling.trainer_run import write_lr_history from ..config.config import TrainingConfig, ValidationConfig from ..config.config_io import save_run_config from ..distributed import grad_tree, launch_distributed_run @@ -237,6 +238,12 @@ def worker( ckpts_path = os.path.join(run_cfg.run_path, "checkpoints") resume = run_cfg.resume + # Record the realized LR schedule beside the checkpoints, in HF's + # log_history shape, so SOURCE can read per-step LRs off a bergson run + # through the same path it already uses for HF Trainer runs. + if global_rank == 0: + write_lr_history(ckpts_path, schedule, len(stream)) + fwd_state = trainer.train( fwd_state, stream, diff --git a/bergson/utils/optimizer_placement.py b/bergson/utils/optimizer_placement.py new file mode 100644 index 00000000..17f202b8 --- /dev/null +++ b/bergson/utils/optimizer_placement.py @@ -0,0 +1,204 @@ +"""Place trainer-written optimizer states where attribution methods look for them. + +The trainer writes each snapshot's second moments as a *sibling* file, +``/step_.optimizer.pt``, next to that step's ``step_.ckpt`` +(see :meth:`bergson.magic.trainer.Trainer.train`). Both attribution consumers +instead expect ``optimizer.pt`` *inside* a directory: + +- SOURCE / approximate unrolling reads ``/optimizer.pt`` for every + checkpoint when ``ApproxUnrollingConfig.use_adam_preconditioner`` is set. +- TrackStar (and any ``build``) resolves ``PreprocessConfig.optimizer_state`` + via :func:`bergson.utils.load_from_optimizer.load_optimizer`, which accepts a + directory and loads ``optimizer.pt`` from inside it. + +Exporting a run to HF-format checkpoint dirs does not carry the sibling files +across, which is what previously forced a manual copy. :func:`place_optimizer_states` +does that placement, and satisfies both consumers with one call since they agree +on the ``/optimizer.pt`` layout. +""" + +import os +import re +import shutil +from pathlib import Path +from typing import Literal, Sequence + +OPTIMIZER_STATE_FILE = "optimizer.pt" +"""Filename both SOURCE and TrackStar look for inside a checkpoint directory.""" + +_STEP_OPTIMIZER_RE = re.compile(r"^step_(\d+)\.optimizer\.pt$") +_STEP_DIR_RE = re.compile(r"step_(\d+)") + +PlacementMode = Literal["copy", "symlink", "hardlink", "move"] + + +def sorted_optimizer_states(save_dir: str | Path) -> list[tuple[int, Path]]: + """Return ``(step, path)`` for each ``step_.optimizer.pt``, sorted by step. + + Mirrors :func:`bergson.data.sorted_checkpoints`, which does the same for the + ``step_.ckpt`` directories these files sit beside. + """ + save_dir = Path(save_dir) + if not save_dir.is_dir(): + raise NotADirectoryError(f"{save_dir} is not a directory") + + found = [] + for name in os.listdir(save_dir): + match = _STEP_OPTIMIZER_RE.match(name) + if match: + found.append((int(match.group(1)), save_dir / name)) + + return sorted(found, key=lambda pair: pair[0]) + + +def _step_of(path: Path) -> int | None: + """Step index named by a checkpoint directory, or None if it names none.""" + match = _STEP_DIR_RE.search(path.name) + return int(match.group(1)) if match else None + + +def _place(src: Path, dst: Path, mode: PlacementMode) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + if dst.exists() or dst.is_symlink(): + dst.unlink() + + match mode: + case "copy": + shutil.copy2(src, dst) + case "move": + shutil.move(str(src), str(dst)) + case "symlink": + # Relative, so an exported run stays valid if the tree is moved. + dst.symlink_to(os.path.relpath(src, dst.parent)) + case "hardlink": + os.link(src, dst) + case other: + raise ValueError(f"Unsupported placement mode: {other}") + + +def place_optimizer_states( + save_dir: str | Path, + checkpoints: Sequence[str | Path], + *, + mode: PlacementMode = "symlink", + overwrite: bool = False, +) -> dict[Path, Path]: + """Put an ``optimizer.pt`` inside each checkpoint dir, from ``save_dir``. + + ``save_dir`` is the trainer's output directory, holding the + ``step_.optimizer.pt`` files written when + ``TrainingConfig.save_optimizer_state`` is set. ``checkpoints`` are the + destination directories -- typically the exported HF checkpoints listed in + ``ApproxUnrollingConfig.checkpoints``. + + States are matched to destinations by step index when the destination + directory names one (``step_12`` / ``step_12.ckpt``); otherwise they are + matched positionally in step order, which requires the counts to be equal. + Mixing the two is an error, since a partial name match usually means the + caller passed a directory list that does not correspond to this run. + + ``mode`` defaults to ``"symlink"``: these files are large (second moments + are one scalar per weight) and duplicating them per checkpoint is wasteful. + Use ``"copy"`` when the destinations must be self-contained, e.g. before + uploading them somewhere. + + Returns ``{destination optimizer.pt: source file}``. Raises rather than + silently skipping, so a wrong directory list fails loudly instead of + producing a preconditioner from the wrong steps. + """ + states = sorted_optimizer_states(save_dir) + if not states: + raise FileNotFoundError( + f"No step_.optimizer.pt files in {save_dir}. Train with " + "TrainingConfig.save_optimizer_state=True to write them." + ) + + dsts = [Path(c) for c in checkpoints] + if not dsts: + raise ValueError("checkpoints is empty; nothing to place") + + missing = [d for d in dsts if not d.is_dir()] + if missing: + raise NotADirectoryError( + f"Checkpoint directories do not exist: {[str(m) for m in missing]}" + ) + + named = [_step_of(d) for d in dsts] + if all(step is not None for step in named): + steps = [step for step in named if step is not None] + by_step = dict(states) + unknown = [step for step in steps if step not in by_step] + if unknown: + raise FileNotFoundError( + f"No step_.optimizer.pt in {save_dir} for steps {unknown}; " + f"available steps: {sorted(by_step)}" + ) + pairs = [(by_step[step], dst) for step, dst in zip(steps, dsts)] + elif any(step is not None for step in named): + raise ValueError( + "Some checkpoint directories name a step and some do not, so they " + "cannot be matched reliably; pass directories that all name their " + "step, or none that do (for positional matching)." + ) + else: + if len(states) != len(dsts): + raise ValueError( + f"Cannot match positionally: {len(states)} optimizer states in " + f"{save_dir} but {len(dsts)} checkpoint directories. Name the " + "destination dirs after their steps to match by step instead." + ) + pairs = [(src, dst) for (_, src), dst in zip(states, dsts)] + + if not overwrite: + existing = [dst / OPTIMIZER_STATE_FILE for _, dst in pairs] + clashes = [p for p in existing if p.exists() or p.is_symlink()] + if clashes: + raise FileExistsError( + f"{[str(c) for c in clashes]} already exist; pass overwrite=True " + "to replace them." + ) + + placed: dict[Path, Path] = {} + for src, dst in pairs: + target = dst / OPTIMIZER_STATE_FILE + _place(src, target, mode) + placed[target] = src + + return placed + + +def place_final_optimizer_state( + run_path: str | Path, + destination: str | Path, + *, + mode: PlacementMode = "symlink", + overwrite: bool = False, +) -> Path: + """Put the run's final ``optimizer.pt`` inside ``destination``. + + ``TrainingConfig.save_optimizer_state`` writes the end-of-training state to + ``/optimizer.pt``. This places it inside an exported model + directory so ``PreprocessConfig.optimizer_state`` can point at that + directory -- the TrackStar gradient-normalization path. + + Returns the path written. + """ + src = Path(run_path) / OPTIMIZER_STATE_FILE + if not src.is_file(): + raise FileNotFoundError( + f"{src} not found. Train with TrainingConfig.save_optimizer_state=True " + "to write it." + ) + + dst_dir = Path(destination) + if not dst_dir.is_dir(): + raise NotADirectoryError(f"{dst_dir} is not a directory") + + target = dst_dir / OPTIMIZER_STATE_FILE + if target.resolve() == src.resolve(): + return target + if not overwrite and (target.exists() or target.is_symlink()): + raise FileExistsError(f"{target} already exists; pass overwrite=True.") + + _place(src, target, mode) + return target diff --git a/bergson/utils/trainer_export.py b/bergson/utils/trainer_export.py new file mode 100644 index 00000000..f9289721 --- /dev/null +++ b/bergson/utils/trainer_export.py @@ -0,0 +1,134 @@ +"""Export a bergson training run's checkpoints to HF-format model directories. + +The trainer saves distributed checkpoints as ``/checkpoints/step_.ckpt`` +(torch.distributed.checkpoint), which SOURCE cannot read -- it loads each +checkpoint with ``from_pretrained``. This turns a run's DCP checkpoints into +``/checkpoint-/`` directories that ``from_pretrained`` accepts, carrying +each step's optimizer state and the run's LR history across so the SOURCE +pipeline finds everything where it expects it. + +Runs offline against a finished run rather than during training, so a run only +pays the extra disk for the trajectories actually being attributed. +""" + +import shutil +from pathlib import Path + +import torch +from transformers import AutoTokenizer + +from ..approx_unrolling.trainer_run import ( + LR_HISTORY_FILENAME, + load_training_config, +) +from ..config.config import TrainingConfig +from ..magic.trainer import prepare_trainer +from .logger import get_logger +from .optimizer_placement import OPTIMIZER_STATE_FILE + +logger = get_logger(__name__) + + +def sorted_dcp_checkpoints(checkpoints_dir: str | Path) -> list[tuple[int, Path]]: + """``(step, path)`` for each ``step_.ckpt``, ordered by step.""" + checkpoints_dir = Path(checkpoints_dir) + if not checkpoints_dir.is_dir(): + raise NotADirectoryError(f"{checkpoints_dir} is not a directory") + + found = [] + for entry in checkpoints_dir.iterdir(): + if entry.is_dir() and entry.name.startswith("step_"): + stem = entry.name.removesuffix(".ckpt").removeprefix("step_") + if stem.isdigit(): + found.append((int(stem), entry)) + + return sorted(found, key=lambda pair: pair[0]) + + +def export_checkpoints( + run_path: str | Path, + out_dir: str | Path | None = None, + *, + training_cfg: TrainingConfig | None = None, + steps: list[int] | None = None, + overwrite: bool = False, +) -> list[Path]: + """Export ``/checkpoints/step_.ckpt`` to ``checkpoint-/`` dirs. + + ``training_cfg`` defaults to the run's own ``config.yaml``, so a finished run + needs no arguments beyond its path. Pass one explicitly to export a run whose + config lives elsewhere. + + ``steps`` restricts the export to those step indices -- useful because SOURCE + needs only ``segments * per_segment`` checkpoints, and each export is a full + model copy. + + Returns the exported directories in step order. + """ + run_path = Path(run_path) + out_dir = Path(out_dir) if out_dir is not None else run_path / "exported" + cfg = training_cfg if training_cfg is not None else load_training_config(run_path) + + available = sorted_dcp_checkpoints(run_path / "checkpoints") + if not available: + raise FileNotFoundError( + f"No step_.ckpt directories under {run_path / 'checkpoints'}; " + "train with a save_mode that writes checkpoints." + ) + + if steps is not None: + by_step = dict(available) + missing = [s for s in steps if s not in by_step] + if missing: + raise FileNotFoundError( + f"Requested steps {missing} were not saved; available: " + f"{[s for s, _ in available]}" + ) + selected = [(s, by_step[s]) for s in sorted(steps)] + else: + selected = available + + out_dir.mkdir(parents=True, exist_ok=True) + existing = [out_dir / f"checkpoint-{s}" for s, _ in selected] + clashes = [p for p in existing if p.exists()] + if clashes and not overwrite: + raise FileExistsError( + f"{[str(c) for c in clashes]} already exist; pass overwrite=True." + ) + + # One model/state is built and reloaded per checkpoint, rather than one per + # step: the state's tensors are loaded in place by TrainerState.load. + trainer, state, model = prepare_trainer(cfg, rank=0, schedule=lambda step: 0.0) + + tokenizer = AutoTokenizer.from_pretrained(cfg.model) + + exported: list[Path] = [] + for step, ckpt in selected: + state.load(str(ckpt)) + + dst = out_dir / f"checkpoint-{step}" + if dst.exists(): + shutil.rmtree(dst) + dst.mkdir(parents=True) + + with state.activate(model), torch.no_grad(): + model.save_pretrained(str(dst), safe_serialization=True) + tokenizer.save_pretrained(str(dst)) + + # SOURCE's Adam variant reads /optimizer.pt; carry the + # trainer's sibling file in under the name it looks for. + sibling = ckpt.parent / f"step_{step}.optimizer.pt" + if sibling.is_file(): + shutil.copy2(sibling, dst / OPTIMIZER_STATE_FILE) + + exported.append(dst) + logger.info("Exported %s -> %s", ckpt.name, dst) + + # The LR history sits beside the checkpoints; SOURCE looks for it beside the + # *exported* ones, since that is the parent of cfg.checkpoints[0]. + history = run_path / "checkpoints" / LR_HISTORY_FILENAME + if history.is_file(): + shutil.copy2(history, out_dir / LR_HISTORY_FILENAME) + + del trainer, state, model + return exported diff --git a/tests/test_optimizer_placement.py b/tests/test_optimizer_placement.py new file mode 100644 index 00000000..bf4534cc --- /dev/null +++ b/tests/test_optimizer_placement.py @@ -0,0 +1,208 @@ +"""Placing trainer-written optimizer states where SOURCE and TrackStar read them. + +The trainer writes ``step_.optimizer.pt`` as a sibling of ``step_.ckpt``; +both consumers want ``optimizer.pt`` *inside* a directory. These tests pin the +matching rules, since a wrong match silently builds a preconditioner from the +wrong step rather than failing. +""" + +import os +import shutil +from pathlib import Path + +import pytest +import torch + +from bergson.utils.load_from_optimizer import load_optimizer +from bergson.utils.optimizer_placement import ( + OPTIMIZER_STATE_FILE, + place_final_optimizer_state, + place_optimizer_states, + sorted_optimizer_states, +) + + +def _write_state(path: Path, step: int) -> None: + """A minimal optimizer.pt whose payload identifies its step.""" + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "state": {0: {"exp_avg_sq": torch.full((2, 3), float(step))}}, + "param_groups": [{"betas": (0.9, 0.999), "eps": 1e-8, "lr": 1e-4}], + }, + path, + ) + + +@pytest.fixture +def save_dir(tmp_path) -> Path: + """A trainer output dir with sibling states at steps 0, 4 and 9.""" + d = tmp_path / "run" + for step in (0, 4, 9): + _write_state(d / f"step_{step}.optimizer.pt", step) + (d / f"step_{step}.ckpt").mkdir(parents=True, exist_ok=True) + return d + + +def _exported(root: Path, names) -> list[Path]: + dirs = [] + for name in names: + p = root / name + p.mkdir(parents=True, exist_ok=True) + dirs.append(p) + return dirs + + +def test_sorted_optimizer_states_orders_numerically(save_dir): + """Steps must sort numerically, not lexically (step_10 < step_9 as strings).""" + _write_state(save_dir / "step_10.optimizer.pt", 10) + assert [s for s, _ in sorted_optimizer_states(save_dir)] == [0, 4, 9, 10] + + +def test_sorted_optimizer_states_ignores_other_files(save_dir): + """The final optimizer.pt and the .ckpt dirs are not per-step states.""" + _write_state(save_dir / OPTIMIZER_STATE_FILE, -1) + assert [s for s, _ in sorted_optimizer_states(save_dir)] == [0, 4, 9] + + +def test_places_by_step_name(save_dir, tmp_path): + """Directories naming a step get that step's state, regardless of order.""" + dsts = _exported(tmp_path / "hf", ["step_9", "step_0", "step_4"]) + placed = place_optimizer_states(save_dir, dsts, mode="copy") + + assert len(placed) == 3 + for dst, expected in zip(dsts, (9, 0, 4)): + blob = load_optimizer(str(dst)) + assert blob["state"][0]["exp_avg_sq"][0, 0].item() == expected + + +def test_places_positionally_when_unnamed(save_dir, tmp_path): + """Unnamed dirs are matched in step order, so seg 0 gets the earliest step.""" + dsts = _exported(tmp_path / "hf", ["a", "b", "c"]) + place_optimizer_states(save_dir, dsts, mode="copy") + + for dst, expected in zip(dsts, (0, 4, 9)): + blob = load_optimizer(str(dst)) + assert blob["state"][0]["exp_avg_sq"][0, 0].item() == expected + + +def test_positional_requires_equal_counts(save_dir, tmp_path): + """A short/long dir list means the caller passed the wrong run; don't guess.""" + dsts = _exported(tmp_path / "hf", ["a", "b"]) + with pytest.raises(ValueError, match="Cannot match positionally"): + place_optimizer_states(save_dir, dsts, mode="copy") + + +def test_mixed_naming_is_rejected(save_dir, tmp_path): + """Half-named lists can't be matched reliably, so refuse rather than guess.""" + dsts = _exported(tmp_path / "hf", ["step_0", "b", "c"]) + with pytest.raises(ValueError, match="name a step and some do not"): + place_optimizer_states(save_dir, dsts, mode="copy") + + +def test_unknown_step_is_rejected(save_dir, tmp_path): + """A named step with no matching state is an error, not a skip.""" + dsts = _exported(tmp_path / "hf", ["step_0", "step_7"]) + with pytest.raises(FileNotFoundError, match=r"steps \[7\]"): + place_optimizer_states(save_dir, dsts, mode="copy") + + +def test_refuses_to_clobber_without_overwrite(save_dir, tmp_path): + """Existing optimizer.pt files are protected; nothing is written on refusal.""" + dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) + (dsts[1] / OPTIMIZER_STATE_FILE).write_text("do not clobber") + + with pytest.raises(FileExistsError): + place_optimizer_states(save_dir, dsts, mode="copy") + + assert (dsts[1] / OPTIMIZER_STATE_FILE).read_text() == "do not clobber" + # The check precedes all writes, so the untouched dirs stay untouched. + assert not (dsts[0] / OPTIMIZER_STATE_FILE).exists() + + place_optimizer_states(save_dir, dsts, mode="copy", overwrite=True) + assert load_optimizer(str(dsts[1]))["state"][0]["exp_avg_sq"][0, 0].item() == 4 + + +def test_symlink_mode_is_relative_and_loadable(save_dir, tmp_path): + """Default mode links rather than duplicating a state per checkpoint.""" + dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) + place_optimizer_states(save_dir, dsts) + + link = dsts[0] / OPTIMIZER_STATE_FILE + assert link.is_symlink() + assert not os.path.isabs(os.readlink(link)), "link must be relative" + assert load_optimizer(str(dsts[0]))["state"][0]["exp_avg_sq"][0, 0].item() == 0 + + +def test_symlinks_survive_relocation(save_dir, tmp_path): + """Relative links keep resolving when run and checkpoints move together.""" + dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) + place_optimizer_states(save_dir, dsts) + + moved = tmp_path / "moved" + moved.mkdir() + shutil.copytree(tmp_path / "hf", moved / "hf", symlinks=True) + shutil.copytree(save_dir, moved / save_dir.name, symlinks=True) + + relocated = moved / "hf" / "step_0" / OPTIMIZER_STATE_FILE + assert relocated.is_symlink() + assert relocated.resolve().is_file(), "relative link broke after relocation" + assert ( + load_optimizer(str(moved / "hf" / "step_0"))["state"][0]["exp_avg_sq"][ + 0, 0 + ].item() + == 0 + ) + + +def test_hardlink_and_move_modes(save_dir, tmp_path): + dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) + place_optimizer_states(save_dir, dsts, mode="hardlink") + assert load_optimizer(str(dsts[2]))["state"][0]["exp_avg_sq"][0, 0].item() == 9 + + dsts2 = _exported(tmp_path / "hf2", ["step_0", "step_4", "step_9"]) + place_optimizer_states(save_dir, dsts2, mode="move") + assert not (save_dir / "step_9.optimizer.pt").exists() + assert load_optimizer(str(dsts2[2]))["state"][0]["exp_avg_sq"][0, 0].item() == 9 + + +def test_missing_states_names_the_config_flag(tmp_path): + """The error should say how to produce the files, not just that they're absent.""" + empty = tmp_path / "empty" + empty.mkdir() + dsts = _exported(tmp_path / "hf", ["step_0"]) + with pytest.raises(FileNotFoundError, match="save_optimizer_state"): + place_optimizer_states(empty, dsts) + + +def test_missing_destination_is_rejected(save_dir, tmp_path): + with pytest.raises(NotADirectoryError): + place_optimizer_states(save_dir, [tmp_path / "nope"]) + + +def test_place_final_state_for_trackstar(save_dir, tmp_path): + """TrackStar points optimizer_state at a dir; the final state must land there.""" + _write_state(save_dir / OPTIMIZER_STATE_FILE, 99) + exported = _exported(tmp_path / "model", ["final"])[0] + + target = place_final_optimizer_state(save_dir, exported, mode="copy") + + assert target == exported / OPTIMIZER_STATE_FILE + # load_optimizer resolves a directory, which is what PreprocessConfig takes. + assert load_optimizer(str(exported))["state"][0]["exp_avg_sq"][0, 0].item() == 99 + + +def test_place_final_state_is_idempotent_in_place(save_dir): + """Pointing it at its own directory is a no-op, not a self-clobber.""" + _write_state(save_dir / OPTIMIZER_STATE_FILE, 99) + target = place_final_optimizer_state(save_dir, save_dir) + assert target.is_file() + assert load_optimizer(str(save_dir))["state"][0]["exp_avg_sq"][0, 0].item() == 99 + + +def test_place_final_state_missing_names_the_flag(tmp_path): + empty = tmp_path / "empty" + empty.mkdir() + dst = _exported(tmp_path / "model", ["final"])[0] + with pytest.raises(FileNotFoundError, match="save_optimizer_state"): + place_final_optimizer_state(empty, dst) diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py new file mode 100644 index 00000000..db481974 --- /dev/null +++ b/tests/test_source_trainer_integration.py @@ -0,0 +1,289 @@ +"""SOURCE reading its hyperparameters off a bergson training run. + +The rule these pin: derivation is a *fallback*. Anything set explicitly wins, so +checkpoints from another trainer keep working exactly as before. +""" + +import json + +import pytest +import torch +import yaml + +from bergson.approx_unrolling.approx_unrolling_math import ( + _checkpoint_step, + compute_lr_times_steps_per_segment, +) +from bergson.approx_unrolling.trainer_run import ( + LR_HISTORY_FILENAME, + derive_momentum, + discover_checkpoints, + load_training_config, + resolve, + write_lr_history, +) +from bergson.config.config import ApproxUnrollingConfig, TrainingConfig + + +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])) + return tmp_path + + +# ── step parsing ──────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "name,expected", + [ + ("checkpoint-120", 120), # HF Trainer + ("step_7.ckpt", 7), # bergson trainer, native + ("step_7", 7), # bergson trainer, exported + ("42", 42), # bare step dir + ], +) +def test_checkpoint_step_accepts_both_conventions(name, expected): + assert _checkpoint_step(f"/runs/x/{name}") == expected + + +def test_checkpoint_step_still_rejects_junk(): + with pytest.raises(ValueError, match="Cannot infer a training step"): + _checkpoint_step("/runs/x/final-model") + + +# ── momentum derivation ───────────────────────────────────────────────────── + + +def test_derive_momentum_sgd_uses_adam_beta1(): + """bergson's SGD passes adam_beta1 as torchopt.sgd's momentum.""" + cfg = TrainingConfig(run_path="/tmp/x", optimizer="sgd", adam_beta1=0.9) + assert derive_momentum(cfg) == 0.9 + + +def test_derive_momentum_sgd_default_is_not_zero(): + """The default adam_beta1 is 0.95, so assuming 0.0 is a 20x lr*steps error.""" + cfg = TrainingConfig(run_path="/tmp/x", optimizer="sgd") + assert derive_momentum(cfg) == pytest.approx(0.95) + assert 1.0 / (1.0 - derive_momentum(cfg)) == pytest.approx(20.0) + + +def test_derive_momentum_adamw_is_zero(): + """AdamW's own preconditioner accounts for its first moment.""" + cfg = TrainingConfig(run_path="/tmp/x", optimizer="adamw", adam_beta1=0.95) + assert derive_momentum(cfg) == 0.0 + + +def test_derive_momentum_muon_warns_and_defaults(caplog): + cfg = TrainingConfig(run_path="/tmp/x", optimizer="muon") + assert derive_momentum(cfg) == 0.0 + + +# ── resolution precedence ─────────────────────────────────────────────────── + + +def test_resolve_is_noop_without_trainer_run(): + """Configs for other trainers pass through untouched.""" + cfg = ApproxUnrollingConfig(checkpoints=["a", "b"], model_path="gpt2") + out = resolve(cfg) + assert out.checkpoints == ["a", "b"] + assert out.model_path == "gpt2" + assert out.momentum == 0.0 # sentinel normalized, nothing derived + + +def test_resolve_fills_momentum_and_model_from_run(tmp_path): + run = _run_dir(tmp_path, optimizer="sgd", adam_beta1=0.9, model="EleutherAI/pythia-14m") + cfg = ApproxUnrollingConfig(trainer_run=str(run), checkpoints=["a", "b"]) + + out = resolve(cfg) + + assert out.momentum == 0.9 + assert out.model_path == "EleutherAI/pythia-14m" + + +def test_explicit_momentum_wins_over_run(tmp_path): + """A user training elsewhere must be able to override what the run says.""" + run = _run_dir(tmp_path, optimizer="sgd", adam_beta1=0.9) + cfg = ApproxUnrollingConfig( + trainer_run=str(run), checkpoints=["a"], momentum=0.5 + ) + assert resolve(cfg).momentum == 0.5 + + +def test_explicit_zero_momentum_is_respected(tmp_path): + """0.0 is a real value, not 'unset' -- it must survive derivation.""" + run = _run_dir(tmp_path, optimizer="sgd", adam_beta1=0.9) + cfg = ApproxUnrollingConfig( + trainer_run=str(run), checkpoints=["a"], momentum=0.0 + ) + assert resolve(cfg).momentum == 0.0 + + +def test_explicit_model_path_wins(tmp_path): + run = _run_dir(tmp_path, model="EleutherAI/pythia-14m") + cfg = ApproxUnrollingConfig( + trainer_run=str(run), checkpoints=["a"], model_path="gpt2" + ) + assert resolve(cfg).model_path == "gpt2" + + +def test_explicit_checkpoints_win(tmp_path): + run = _run_dir(tmp_path) + (run / "checkpoint-0").mkdir() + cfg = ApproxUnrollingConfig(trainer_run=str(run), checkpoints=["mine"]) + assert resolve(cfg).checkpoints == ["mine"] + + +def test_resolve_discovers_exported_checkpoints_in_order(tmp_path): + run = _run_dir(tmp_path) + for step in (0, 10, 2): + (run / f"checkpoint-{step}").mkdir() + + out = resolve(ApproxUnrollingConfig(trainer_run=str(run))) + + assert [p.split("-")[-1] for p in out.checkpoints] == ["0", "2", "10"] + + +def test_unexported_run_names_the_export_step(tmp_path): + """Native DCP checkpoints can't be loaded by from_pretrained; say so.""" + run = _run_dir(tmp_path) + (run / "checkpoints").mkdir() + (run / "checkpoints" / "step_0.ckpt").mkdir() + + with pytest.raises(FileNotFoundError, match="export_checkpoints"): + resolve(ApproxUnrollingConfig(trainer_run=str(run))) + + +def test_missing_config_yaml_is_explained(tmp_path): + with pytest.raises(FileNotFoundError, match="bergson run directory"): + load_training_config(tmp_path) + + +# ── LR history ────────────────────────────────────────────────────────────── + + +def test_write_lr_history_matches_hf_log_history_shape(tmp_path): + """Written in HF's shape so SOURCE's existing reader picks it up unchanged.""" + path = write_lr_history(tmp_path, lambda step: 1e-4 * (step + 1), 3) + + assert path.name == LR_HISTORY_FILENAME + entries = json.loads(path.read_text()) + assert entries == [ + {"step": 0, "learning_rate": pytest.approx(1e-4)}, + {"step": 1, "learning_rate": pytest.approx(2e-4)}, + {"step": 2, "learning_rate": pytest.approx(3e-4)}, + ] + + +def test_lr_times_steps_reads_bergson_history(tmp_path): + """End to end: a written history drives lr*K without any HF artifacts.""" + export = tmp_path / "exported" + export.mkdir() + for step in (2, 4): + (export / f"checkpoint-{step}").mkdir() + write_lr_history(export, lambda step: 1e-3, 5) + + cfg = ApproxUnrollingConfig( + checkpoints=[str(export / "checkpoint-2"), str(export / "checkpoint-4")], + segments=2, + momentum=0.0, + ) + + # Segment 1 covers steps 1..2, segment 2 covers 3..4 -> two steps each. + assert compute_lr_times_steps_per_segment(cfg) == [ + pytest.approx(2e-3), + pytest.approx(2e-3), + ] + + +def test_momentum_scales_lr_times_steps(tmp_path): + """The 1/(1-beta) terminal-velocity factor is what the SGD fix is about.""" + cfg = ApproxUnrollingConfig( + checkpoints=["a", "b"], + segments=2, + lr_list=[1e-3, 1e-3], + step_size_list=[10, 10], + momentum=0.0, + ) + baseline = compute_lr_times_steps_per_segment(cfg) + + cfg.momentum = 0.95 + scaled = compute_lr_times_steps_per_segment(cfg) + + assert scaled == [pytest.approx(20 * b) for b in baseline] + + +def test_momentum_out_of_range_is_rejected(): + cfg = ApproxUnrollingConfig( + checkpoints=["a"], segments=1, lr_list=[1e-3], step_size_list=[1], momentum=1.0 + ) + with pytest.raises(ValueError, match="momentum must be in"): + compute_lr_times_steps_per_segment(cfg) + + +def test_unset_momentum_defaults_to_no_scaling(): + """Without a trainer_run, behaviour is exactly as before this change.""" + cfg = ApproxUnrollingConfig( + checkpoints=["a"], segments=1, lr_list=[1e-3], step_size_list=[10] + ) + assert compute_lr_times_steps_per_segment(cfg) == [pytest.approx(1e-2)] + + +# ── export end to end ─────────────────────────────────────────────────────── + + +def test_export_round_trips_checkpoint_weights(tmp_path): + """A DCP checkpoint must survive export as a from_pretrained-loadable model. + + SOURCE loads every checkpoint with from_pretrained, so an export that lost + or mangled weights would silently attribute the wrong trajectory. + """ + import torchopt + from datasets import Dataset + from transformers import AutoConfig, AutoModelForCausalLM + + from bergson.magic.data_stream import DataStream + from bergson.magic.trainer import Trainer + from bergson.utils.trainer_export import sorted_dcp_checkpoints + + torch.manual_seed(0) + config = AutoConfig.from_pretrained("EleutherAI/pythia-14m") + + def fresh(): + torch.manual_seed(0) + m = AutoModelForCausalLM.from_config( + config, dtype=torch.float32, attn_implementation="eager" + ) + m.requires_grad_(True) + return m + + n = 4 + ds = Dataset.from_dict( + {"input_ids": [[1, 2, 3, 4]] * n, "labels": [[1, 2, 3, 4]] * n} + ) + stream = DataStream(ds, batch_size=1, device="cpu") + opt = torchopt.sgd(lambda step: 1e-4, momentum=0.95) + + trainer, state = Trainer.initialize(fresh(), opt) + save_dir = tmp_path / "checkpoints" + trainer.train(state, stream, inplace=True, save_dir=str(save_dir), save_mode="all") + + found = sorted_dcp_checkpoints(save_dir) + assert [s for s, _ in found] == list(range(n)) + + # Reload the last checkpoint and export it, as export_checkpoints does. + model = fresh() + _, loaded = Trainer.initialize(model, opt) + loaded.load(str(found[-1][1])) + + out = tmp_path / "checkpoint-3" + with loaded.activate(model), torch.no_grad(): + model.save_pretrained(str(out), safe_serialization=True) + reference = {k: v.detach().clone() for k, v in model.named_parameters()} + + reloaded = AutoModelForCausalLM.from_pretrained(str(out)) + got = dict(reloaded.named_parameters()) + for name, ref in reference.items(): + torch.testing.assert_close(got[name], ref, atol=0, rtol=0) From da08ff2a21fd30e99069a0b628e2fe21647ae501 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 12:43:32 +0000 Subject: [PATCH 02/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_source_trainer_integration.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index db481974..e02b7cfc 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -17,7 +17,6 @@ from bergson.approx_unrolling.trainer_run import ( LR_HISTORY_FILENAME, derive_momentum, - discover_checkpoints, load_training_config, resolve, write_lr_history, @@ -94,7 +93,9 @@ def test_resolve_is_noop_without_trainer_run(): def test_resolve_fills_momentum_and_model_from_run(tmp_path): - run = _run_dir(tmp_path, optimizer="sgd", adam_beta1=0.9, model="EleutherAI/pythia-14m") + run = _run_dir( + tmp_path, optimizer="sgd", adam_beta1=0.9, model="EleutherAI/pythia-14m" + ) cfg = ApproxUnrollingConfig(trainer_run=str(run), checkpoints=["a", "b"]) out = resolve(cfg) @@ -106,18 +107,14 @@ def test_resolve_fills_momentum_and_model_from_run(tmp_path): def test_explicit_momentum_wins_over_run(tmp_path): """A user training elsewhere must be able to override what the run says.""" run = _run_dir(tmp_path, optimizer="sgd", adam_beta1=0.9) - cfg = ApproxUnrollingConfig( - trainer_run=str(run), checkpoints=["a"], momentum=0.5 - ) + cfg = ApproxUnrollingConfig(trainer_run=str(run), checkpoints=["a"], momentum=0.5) assert resolve(cfg).momentum == 0.5 def test_explicit_zero_momentum_is_respected(tmp_path): """0.0 is a real value, not 'unset' -- it must survive derivation.""" run = _run_dir(tmp_path, optimizer="sgd", adam_beta1=0.9) - cfg = ApproxUnrollingConfig( - trainer_run=str(run), checkpoints=["a"], momentum=0.0 - ) + cfg = ApproxUnrollingConfig(trainer_run=str(run), checkpoints=["a"], momentum=0.0) assert resolve(cfg).momentum == 0.0 From 0d360c2158b4c98847b40aad83797c60b1b43d56 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 13:02:12 +0000 Subject: [PATCH 03/13] refactor: drop the import alias, trim docstrings --- .../approx_unrolling/approx_unrolling_math.py | 13 ++-- bergson/approx_unrolling/pipeline.py | 8 +-- bergson/approx_unrolling/trainer_run.py | 65 +++++-------------- bergson/config/config.py | 22 ++----- bergson/utils/trainer_export.py | 27 ++------ 5 files changed, 37 insertions(+), 98 deletions(-) diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index 21ace167..d0445d27 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -29,8 +29,8 @@ def _checkpoint_step(p: str) -> int: """Extract the training step index for a checkpoint. Accepts a path whose final component is a ``checkpoint-`` directory - (what HF Trainer writes natively), a ``step_[.ckpt]`` directory (what - the bergson trainer writes), or a bare ```` step directory. + (HF Trainer), a ``step_[.ckpt]`` directory (Bergson trainer), or a + bare ```` step directory. """ # TODO: Inferring the step from a checkpoint name via regex is brittle. # For now, we raise an error @@ -46,7 +46,7 @@ def _checkpoint_step(p: str) -> int: raise ValueError( f"Cannot infer a training step from checkpoint {p!r}: expected a path " "ending in 'checkpoint-', 'step_[.ckpt]', or a bare '' step " - "directory. Set both " + "directory. Set " "`lr_list` and `step_size_list` on ApproxUnrollingConfig to specify " "the per-segment learning rate and step counts explicitly instead of " "inferring them from checkpoint names." @@ -54,17 +54,14 @@ def _checkpoint_step(p: str) -> int: def compute_lr_times_steps_per_segment( - approx_unrolling_cfg: ApproxUnrollingConfig, + cfg: ApproxUnrollingConfig, ) -> list[float]: """Per-segment lr * K. Use ``lr_list * step_size_list`` if set on config; else equal-partition log_history.json into segments and sum per-step LRs. With SGD heavy-ball ``momentum`` beta, scale by the terminal velocity - 1/(1-beta) (Bae et al. 2024, App. D.2). ``momentum`` is taken as given when - set; ``None`` means it is derived from ``trainer_run`` (or 0.0 without one), - so checkpoints from another trainer are unaffected by the derivation. + 1/(1-beta) (Bae et al. 2024, App. D.2). """ - cfg = approx_unrolling_cfg L = cfg.segments momentum = cfg.momentum if cfg.momentum is not None else 0.0 diff --git a/bergson/approx_unrolling/pipeline.py b/bergson/approx_unrolling/pipeline.py index 7a179cb8..a2853c8b 100644 --- a/bergson/approx_unrolling/pipeline.py +++ b/bergson/approx_unrolling/pipeline.py @@ -52,7 +52,7 @@ aggregate_segment_covariances, aggregate_segment_lambdas, ) -from .trainer_run import resolve as resolve_trainer_run +from .trainer_run import resolve # Total number of steps in the full approximate unrolling pipeline. Used only for the # user-facing "Step k/N_TOTAL_STEPS:" prefix. Bump as steps land. @@ -83,10 +83,8 @@ def approx_unrolling_pipeline( """ logger = get_logger("approx_unrolling_pipeline") - # Fill in anything the config left unset from the bergson run it names. - # A no-op without `trainer_run`, and it never overrides an explicit value, - # so configs for checkpoints from other trainers are unaffected. - approx_unrolling_cfg = resolve_trainer_run(approx_unrolling_cfg) + # Use the Bergson training run to resolve unset config fields if present. + approx_unrolling_cfg = resolve(approx_unrolling_cfg) n_ckpts = len(approx_unrolling_cfg.checkpoints) n_segments = approx_unrolling_cfg.segments diff --git a/bergson/approx_unrolling/trainer_run.py b/bergson/approx_unrolling/trainer_run.py index 740874b1..9906b03e 100644 --- a/bergson/approx_unrolling/trainer_run.py +++ b/bergson/approx_unrolling/trainer_run.py @@ -1,15 +1,7 @@ -"""Fill in SOURCE hyperparameters from a bergson training run. - -SOURCE was built around HF Trainer runs: it infers per-step learning rates from -``log_history.json`` / ``trainer_state.json`` and reads steps out of -``checkpoint-`` directory names. A bergson run has the same information in -its own form -- ``config.yaml`` plus the checkpoint schedule -- so this module -maps one onto the other. - -Everything here is a *fallback*. An explicitly configured field always wins, so -checkpoints produced by any other trainer keep working exactly as before by -setting ``lr_list``/``step_size_list``/``momentum`` by hand and leaving -``trainer_run`` empty. +"""Fill in SOURCE hyperparameters from a bergson run's ``config.yaml``. + +SOURCE was built for HF Trainer runs. Everything here is a fallback: explicit +config fields always win, so other trainers are unaffected. """ import json @@ -24,13 +16,8 @@ from ..utils.logger import get_logger LR_HISTORY_FILENAME = "log_history.json" -"""Per-step learning rates, in HF Trainer's ``log_history`` shape. - -Written next to a bergson run's checkpoints so -:func:`~bergson.approx_unrolling.approx_unrolling_math.compute_lr_times_steps_per_segment` -picks it up through the path it already checks first, with no bergson-specific -branch in the LR math. -""" +"""Per-step LRs in HF's ``log_history`` shape, written beside a run's +checkpoints -- the path the LR math already checks first.""" logger = get_logger(__name__) @@ -38,11 +25,8 @@ def write_lr_history( save_dir: str | Path, schedule: Callable[[int], float], num_steps: int ) -> Path: - """Record the run's realized per-step learning rates beside its checkpoints. - - ``schedule`` is the step -> lr callable the optimizer was built with, so - this is the exact schedule that ran rather than a reconstruction of it. - """ + """Record per-step LRs beside the checkpoints, from the ``schedule`` the + optimizer was built with, so it is exact rather than reconstructed.""" save_dir = Path(save_dir) save_dir.mkdir(parents=True, exist_ok=True) history = [ @@ -56,11 +40,8 @@ def write_lr_history( def load_training_config(trainer_run: str | Path) -> TrainingConfig: - """Load the ``TrainingConfig`` a bergson run was launched with. - - ``save_run_config`` writes ``config.yaml`` as ``{command_name: {...}}``, so - the single value is the run's config regardless of which command wrote it. - """ + """Load the ``TrainingConfig`` a run was launched with. ``save_run_config`` + writes ``{command_name: {...}}``, so the single value is the config.""" path = Path(trainer_run) / CONFIG_FILENAME if not path.is_file(): raise FileNotFoundError( @@ -87,13 +68,8 @@ def load_training_config(trainer_run: str | Path) -> TrainingConfig: def derive_momentum(training_cfg: TrainingConfig) -> float: - """Heavy-ball momentum beta that a bergson run actually trained with. - - bergson's SGD passes ``adam_beta1`` as ``torchopt.sgd``'s momentum (see - ``prepare_trainer``), so the SOURCE-relevant beta is ``adam_beta1`` for SGD - and ``0.0`` for AdamW, whose own preconditioner already accounts for its - first moment. - """ + """Momentum beta a run trained with. bergson's SGD passes ``adam_beta1`` to + ``torchopt.sgd``; AdamW's own preconditioner handles its first moment.""" match training_cfg.optimizer: case "sgd": return float(training_cfg.adam_beta1) @@ -111,13 +87,8 @@ def derive_momentum(training_cfg: TrainingConfig) -> float: def discover_checkpoints(trainer_run: str | Path) -> list[str]: - """The run's saved checkpoints, in training order. - - Prefers exported HF-format ``checkpoint-`` directories (what SOURCE can - actually load) and falls back to the trainer's native ``step_.ckpt``, so - the error surfaced to the caller names the export step rather than an empty - list. - """ + """Saved checkpoints in training order. Prefers exported ``checkpoint-`` + dirs; finding only native ``step_.ckpt`` raises, naming the export.""" root = Path(trainer_run) exported = sorted( (p for p in root.glob("checkpoint-*") if p.is_dir()), @@ -147,12 +118,8 @@ def discover_checkpoints(trainer_run: str | Path) -> list[str]: def resolve(cfg: ApproxUnrollingConfig) -> ApproxUnrollingConfig: - """Return ``cfg`` with unset fields filled in from ``cfg.trainer_run``. - - A no-op when ``trainer_run`` is empty, so configs for other trainers pass - through untouched. Fields already set are never overwritten -- this only - ever supplies a value the caller did not give. - """ + """Fill unset fields from ``cfg.trainer_run``. A no-op when it is empty; + never overwrites a field the caller set.""" if not cfg.trainer_run: # Nothing to derive from; only normalize the momentum sentinel. if cfg.momentum is None: diff --git a/bergson/config/config.py b/bergson/config/config.py index e21de0c0..7cdf43f3 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -787,23 +787,15 @@ class ApproxUnrollingConfig(Serializable): empty, inferred from per-checkpoint step counts.""" trainer_run: str = "" - """Directory of a bergson training run, used to fill in what this config - does not state explicitly: ``checkpoints``, ``model_path`` and - ``momentum``. Leave empty for checkpoints from another trainer (e.g. HF - Trainer) and set those fields by hand -- anything set explicitly always - wins over what is read from the run.""" + """Bergson run directory to read ``checkpoints``, ``model_path`` and + ``momentum`` from. Explicit fields always win; leave empty for other + trainers.""" momentum: float | None = None - """SGD heavy-ball momentum beta used during training; scales each segment's - lr*steps by the terminal velocity 1/(1-beta) (Bae et al. 2024, App. D.2). - - ``None`` means "not set": derive it from ``trainer_run`` if one is given, - else fall back to ``0.0``. Set it explicitly (``0.0`` included) when - training happened outside bergson. - - Deriving matters because bergson's SGD takes its momentum from - ``TrainingConfig.adam_beta1``, which defaults to ``0.95`` -- assuming 0.0 - there understates lr*steps by 20x.""" + """SGD heavy-ball momentum beta; scales lr*steps by 1/(1-beta) (Bae et al. + 2024, App. D.2). ``None`` derives it from ``trainer_run``, else 0.0. + bergson's SGD uses ``adam_beta1`` (default 0.95), so assuming 0.0 there + understates lr*steps by 20x.""" query: DataConfig = field(default_factory=DataConfig) """Query dataset spec; gradients computed at the final checkpoint.""" diff --git a/bergson/utils/trainer_export.py b/bergson/utils/trainer_export.py index f9289721..c6bd1eab 100644 --- a/bergson/utils/trainer_export.py +++ b/bergson/utils/trainer_export.py @@ -1,14 +1,6 @@ -"""Export a bergson training run's checkpoints to HF-format model directories. - -The trainer saves distributed checkpoints as ``/checkpoints/step_.ckpt`` -(torch.distributed.checkpoint), which SOURCE cannot read -- it loads each -checkpoint with ``from_pretrained``. This turns a run's DCP checkpoints into -``/checkpoint-/`` directories that ``from_pretrained`` accepts, carrying -each step's optimizer state and the run's LR history across so the SOURCE -pipeline finds everything where it expects it. - -Runs offline against a finished run rather than during training, so a run only -pays the extra disk for the trajectories actually being attributed. +"""Export a run's DCP ``step_.ckpt`` checkpoints to ``checkpoint-/`` model +dirs that ``from_pretrained`` can load, carrying optimizer state and LR history +across. Offline, so a run only pays disk for trajectories being attributed. """ import shutil @@ -53,17 +45,10 @@ def export_checkpoints( steps: list[int] | None = None, overwrite: bool = False, ) -> list[Path]: - """Export ``/checkpoints/step_.ckpt`` to ``checkpoint-/`` dirs. - - ``training_cfg`` defaults to the run's own ``config.yaml``, so a finished run - needs no arguments beyond its path. Pass one explicitly to export a run whose - config lives elsewhere. - - ``steps`` restricts the export to those step indices -- useful because SOURCE - needs only ``segments * per_segment`` checkpoints, and each export is a full - model copy. + """Export ``step_.ckpt`` to ``checkpoint-/`` dirs, in step order. - Returns the exported directories in step order. + ``training_cfg`` defaults to the run's ``config.yaml``. ``steps`` limits the + export, since each one is a full model copy. """ run_path = Path(run_path) out_dir = Path(out_dir) if out_dir is not None else run_path / "exported" From 337b93f6a81f51a2e0755fd46183203e3619e01b Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 13:08:18 +0000 Subject: [PATCH 04/13] refactor: read the LR history from the run that wrote it The trainer wrote log_history.json beside its checkpoints and the export copied it next to the exported ones, because the reader derived the path from checkpoints[0].parent. Two locations that could drift, and a comment to explain why. Resolve it through trainer_run instead, so the file has one home and the export copies nothing. --- .../approx_unrolling/approx_unrolling_math.py | 12 ++++++---- bergson/utils/trainer_export.py | 17 +++---------- tests/test_source_trainer_integration.py | 24 +++++++++++++++++++ 3 files changed, 34 insertions(+), 19 deletions(-) diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index d0445d27..411bdc31 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -24,6 +24,8 @@ from bergson.score.score import score_dataset from bergson.score.score_writer import save_sequence_scores, save_token_scores +from .trainer_run import LR_HISTORY_FILENAME, lr_history_path + def _checkpoint_step(p: str) -> int: """Extract the training step index for a checkpoint. @@ -77,11 +79,11 @@ def compute_lr_times_steps_per_segment( per_segment = len(cfg.checkpoints) // L ckpt_steps = [_checkpoint_step(p) for p in cfg.checkpoints] boundaries = [0] + [ckpt_steps[(l + 1) * per_segment - 1] for l in range(L)] - # Prefer log_history.json if dumped at the parent dir; otherwise pull - # log_history out of the final checkpoint's trainer_state.json (what HF - # Trainer writes natively). - parent = Path(str(cfg.checkpoints[0])).parent - log_path = parent / "log_history.json" + # A bergson run's own history, else one dumped beside the checkpoints, + # else the final checkpoint's trainer_state.json (what HF Trainer writes). + log_path = lr_history_path(cfg) or ( + Path(str(cfg.checkpoints[0])).parent / LR_HISTORY_FILENAME + ) if log_path.exists(): with open(log_path) as f: log_history = json.load(f) diff --git a/bergson/utils/trainer_export.py b/bergson/utils/trainer_export.py index c6bd1eab..a8abb1b0 100644 --- a/bergson/utils/trainer_export.py +++ b/bergson/utils/trainer_export.py @@ -1,7 +1,5 @@ -"""Export a run's DCP ``step_.ckpt`` checkpoints to ``checkpoint-/`` model -dirs that ``from_pretrained`` can load, carrying optimizer state and LR history -across. Offline, so a run only pays disk for trajectories being attributed. -""" +"""Export a run's DCP ``step_.ckpt`` checkpoints to HF-compatible +``checkpoint-/`` model dirs.""" import shutil from pathlib import Path @@ -9,10 +7,7 @@ import torch from transformers import AutoTokenizer -from ..approx_unrolling.trainer_run import ( - LR_HISTORY_FILENAME, - load_training_config, -) +from ..approx_unrolling.trainer_run import load_training_config from ..config.config import TrainingConfig from ..magic.trainer import prepare_trainer from .logger import get_logger @@ -109,11 +104,5 @@ def export_checkpoints( exported.append(dst) logger.info("Exported %s -> %s", ckpt.name, dst) - # The LR history sits beside the checkpoints; SOURCE looks for it beside the - # *exported* ones, since that is the parent of cfg.checkpoints[0]. - history = run_path / "checkpoints" / LR_HISTORY_FILENAME - if history.is_file(): - shutil.copy2(history, out_dir / LR_HISTORY_FILENAME) - del trainer, state, model return exported diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index e02b7cfc..8c75e49b 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -284,3 +284,27 @@ def fresh(): got = dict(reloaded.named_parameters()) for name, ref in reference.items(): torch.testing.assert_close(got[name], ref, atol=0, rtol=0) + + +def test_lr_history_read_from_the_run_not_the_export(tmp_path): + """One logical location: the trainer writes it beside its own checkpoints + and the reader finds it there, so nothing has to be copied on export.""" + run = _run_dir(tmp_path) + write_lr_history(run / "checkpoints", lambda step: 1e-3, 5) + + export = tmp_path / "exported" + export.mkdir() + for step in (2, 4): + (export / f"checkpoint-{step}").mkdir() + assert not (export / LR_HISTORY_FILENAME).exists() + + cfg = ApproxUnrollingConfig( + trainer_run=str(run), + checkpoints=[str(export / "checkpoint-2"), str(export / "checkpoint-4")], + segments=2, + momentum=0.0, + ) + assert compute_lr_times_steps_per_segment(cfg) == [ + pytest.approx(2e-3), + pytest.approx(2e-3), + ] From 5f1db9363dad533bffa9f8295c0e57f3cbe8679a Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 13:25:30 +0000 Subject: [PATCH 05/13] refactor: keep each checkpoint's optimizer state inside the checkpoint The trainer wrote optimizer state as a sibling of the checkpoint it belonged to, which severed the association and forced everything downstream to rebuild it: optimizer_placement.py matched loose files back to checkpoint dirs by step name, or positionally when the names did not say, and had to reject half-named lists because a wrong match would silently build a preconditioner from another step's moments. Writing it inside step_.ckpt/ removes the problem rather than solving it. The file travels with its checkpoint, so the export is a plain copy from one directory into another and nothing needs matching. Verified DCP is unbothered: an extra non-DCP file inside a checkpoint dir leaves dcp.load working, a resumed run reproducing the same parameters, and the file itself intact across a re-save -- pinned by a test, since the layout depends on it. place_final_optimizer_state went too. TrackStar's load_optimizer already takes either a file or a directory, so pointing optimizer_state at /optimizer.pt needs no placement step at all; that helper solved a problem that did not exist. No migration path for the old sibling layout, by choice: re-export instead. --- bergson/utils/optimizer_placement.py | 204 ---------------------- bergson/utils/trainer_export.py | 14 +- tests/test_optimizer_placement.py | 208 ----------------------- tests/test_source_trainer_integration.py | 59 +++++++ 4 files changed, 67 insertions(+), 418 deletions(-) delete mode 100644 bergson/utils/optimizer_placement.py delete mode 100644 tests/test_optimizer_placement.py diff --git a/bergson/utils/optimizer_placement.py b/bergson/utils/optimizer_placement.py deleted file mode 100644 index 17f202b8..00000000 --- a/bergson/utils/optimizer_placement.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Place trainer-written optimizer states where attribution methods look for them. - -The trainer writes each snapshot's second moments as a *sibling* file, -``/step_.optimizer.pt``, next to that step's ``step_.ckpt`` -(see :meth:`bergson.magic.trainer.Trainer.train`). Both attribution consumers -instead expect ``optimizer.pt`` *inside* a directory: - -- SOURCE / approximate unrolling reads ``/optimizer.pt`` for every - checkpoint when ``ApproxUnrollingConfig.use_adam_preconditioner`` is set. -- TrackStar (and any ``build``) resolves ``PreprocessConfig.optimizer_state`` - via :func:`bergson.utils.load_from_optimizer.load_optimizer`, which accepts a - directory and loads ``optimizer.pt`` from inside it. - -Exporting a run to HF-format checkpoint dirs does not carry the sibling files -across, which is what previously forced a manual copy. :func:`place_optimizer_states` -does that placement, and satisfies both consumers with one call since they agree -on the ``/optimizer.pt`` layout. -""" - -import os -import re -import shutil -from pathlib import Path -from typing import Literal, Sequence - -OPTIMIZER_STATE_FILE = "optimizer.pt" -"""Filename both SOURCE and TrackStar look for inside a checkpoint directory.""" - -_STEP_OPTIMIZER_RE = re.compile(r"^step_(\d+)\.optimizer\.pt$") -_STEP_DIR_RE = re.compile(r"step_(\d+)") - -PlacementMode = Literal["copy", "symlink", "hardlink", "move"] - - -def sorted_optimizer_states(save_dir: str | Path) -> list[tuple[int, Path]]: - """Return ``(step, path)`` for each ``step_.optimizer.pt``, sorted by step. - - Mirrors :func:`bergson.data.sorted_checkpoints`, which does the same for the - ``step_.ckpt`` directories these files sit beside. - """ - save_dir = Path(save_dir) - if not save_dir.is_dir(): - raise NotADirectoryError(f"{save_dir} is not a directory") - - found = [] - for name in os.listdir(save_dir): - match = _STEP_OPTIMIZER_RE.match(name) - if match: - found.append((int(match.group(1)), save_dir / name)) - - return sorted(found, key=lambda pair: pair[0]) - - -def _step_of(path: Path) -> int | None: - """Step index named by a checkpoint directory, or None if it names none.""" - match = _STEP_DIR_RE.search(path.name) - return int(match.group(1)) if match else None - - -def _place(src: Path, dst: Path, mode: PlacementMode) -> None: - dst.parent.mkdir(parents=True, exist_ok=True) - if dst.exists() or dst.is_symlink(): - dst.unlink() - - match mode: - case "copy": - shutil.copy2(src, dst) - case "move": - shutil.move(str(src), str(dst)) - case "symlink": - # Relative, so an exported run stays valid if the tree is moved. - dst.symlink_to(os.path.relpath(src, dst.parent)) - case "hardlink": - os.link(src, dst) - case other: - raise ValueError(f"Unsupported placement mode: {other}") - - -def place_optimizer_states( - save_dir: str | Path, - checkpoints: Sequence[str | Path], - *, - mode: PlacementMode = "symlink", - overwrite: bool = False, -) -> dict[Path, Path]: - """Put an ``optimizer.pt`` inside each checkpoint dir, from ``save_dir``. - - ``save_dir`` is the trainer's output directory, holding the - ``step_.optimizer.pt`` files written when - ``TrainingConfig.save_optimizer_state`` is set. ``checkpoints`` are the - destination directories -- typically the exported HF checkpoints listed in - ``ApproxUnrollingConfig.checkpoints``. - - States are matched to destinations by step index when the destination - directory names one (``step_12`` / ``step_12.ckpt``); otherwise they are - matched positionally in step order, which requires the counts to be equal. - Mixing the two is an error, since a partial name match usually means the - caller passed a directory list that does not correspond to this run. - - ``mode`` defaults to ``"symlink"``: these files are large (second moments - are one scalar per weight) and duplicating them per checkpoint is wasteful. - Use ``"copy"`` when the destinations must be self-contained, e.g. before - uploading them somewhere. - - Returns ``{destination optimizer.pt: source file}``. Raises rather than - silently skipping, so a wrong directory list fails loudly instead of - producing a preconditioner from the wrong steps. - """ - states = sorted_optimizer_states(save_dir) - if not states: - raise FileNotFoundError( - f"No step_.optimizer.pt files in {save_dir}. Train with " - "TrainingConfig.save_optimizer_state=True to write them." - ) - - dsts = [Path(c) for c in checkpoints] - if not dsts: - raise ValueError("checkpoints is empty; nothing to place") - - missing = [d for d in dsts if not d.is_dir()] - if missing: - raise NotADirectoryError( - f"Checkpoint directories do not exist: {[str(m) for m in missing]}" - ) - - named = [_step_of(d) for d in dsts] - if all(step is not None for step in named): - steps = [step for step in named if step is not None] - by_step = dict(states) - unknown = [step for step in steps if step not in by_step] - if unknown: - raise FileNotFoundError( - f"No step_.optimizer.pt in {save_dir} for steps {unknown}; " - f"available steps: {sorted(by_step)}" - ) - pairs = [(by_step[step], dst) for step, dst in zip(steps, dsts)] - elif any(step is not None for step in named): - raise ValueError( - "Some checkpoint directories name a step and some do not, so they " - "cannot be matched reliably; pass directories that all name their " - "step, or none that do (for positional matching)." - ) - else: - if len(states) != len(dsts): - raise ValueError( - f"Cannot match positionally: {len(states)} optimizer states in " - f"{save_dir} but {len(dsts)} checkpoint directories. Name the " - "destination dirs after their steps to match by step instead." - ) - pairs = [(src, dst) for (_, src), dst in zip(states, dsts)] - - if not overwrite: - existing = [dst / OPTIMIZER_STATE_FILE for _, dst in pairs] - clashes = [p for p in existing if p.exists() or p.is_symlink()] - if clashes: - raise FileExistsError( - f"{[str(c) for c in clashes]} already exist; pass overwrite=True " - "to replace them." - ) - - placed: dict[Path, Path] = {} - for src, dst in pairs: - target = dst / OPTIMIZER_STATE_FILE - _place(src, target, mode) - placed[target] = src - - return placed - - -def place_final_optimizer_state( - run_path: str | Path, - destination: str | Path, - *, - mode: PlacementMode = "symlink", - overwrite: bool = False, -) -> Path: - """Put the run's final ``optimizer.pt`` inside ``destination``. - - ``TrainingConfig.save_optimizer_state`` writes the end-of-training state to - ``/optimizer.pt``. This places it inside an exported model - directory so ``PreprocessConfig.optimizer_state`` can point at that - directory -- the TrackStar gradient-normalization path. - - Returns the path written. - """ - src = Path(run_path) / OPTIMIZER_STATE_FILE - if not src.is_file(): - raise FileNotFoundError( - f"{src} not found. Train with TrainingConfig.save_optimizer_state=True " - "to write it." - ) - - dst_dir = Path(destination) - if not dst_dir.is_dir(): - raise NotADirectoryError(f"{dst_dir} is not a directory") - - target = dst_dir / OPTIMIZER_STATE_FILE - if target.resolve() == src.resolve(): - return target - if not overwrite and (target.exists() or target.is_symlink()): - raise FileExistsError(f"{target} already exists; pass overwrite=True.") - - _place(src, target, mode) - return target diff --git a/bergson/utils/trainer_export.py b/bergson/utils/trainer_export.py index a8abb1b0..f4140eab 100644 --- a/bergson/utils/trainer_export.py +++ b/bergson/utils/trainer_export.py @@ -11,7 +11,11 @@ from ..config.config import TrainingConfig from ..magic.trainer import prepare_trainer from .logger import get_logger -from .optimizer_placement import OPTIMIZER_STATE_FILE + +OPTIMIZER_STATE_FILE = "optimizer.pt" +"""Second moments, written inside the checkpoint dir they belong to -- both by +the trainer and here, so the file travels with its checkpoint and never has to +be matched back to one.""" logger = get_logger(__name__) @@ -95,11 +99,9 @@ def export_checkpoints( model.save_pretrained(str(dst), safe_serialization=True) tokenizer.save_pretrained(str(dst)) - # SOURCE's Adam variant reads /optimizer.pt; carry the - # trainer's sibling file in under the name it looks for. - sibling = ckpt.parent / f"step_{step}.optimizer.pt" - if sibling.is_file(): - shutil.copy2(sibling, dst / OPTIMIZER_STATE_FILE) + state_file = ckpt / OPTIMIZER_STATE_FILE + if state_file.is_file(): + shutil.copy2(state_file, dst / OPTIMIZER_STATE_FILE) exported.append(dst) logger.info("Exported %s -> %s", ckpt.name, dst) diff --git a/tests/test_optimizer_placement.py b/tests/test_optimizer_placement.py deleted file mode 100644 index bf4534cc..00000000 --- a/tests/test_optimizer_placement.py +++ /dev/null @@ -1,208 +0,0 @@ -"""Placing trainer-written optimizer states where SOURCE and TrackStar read them. - -The trainer writes ``step_.optimizer.pt`` as a sibling of ``step_.ckpt``; -both consumers want ``optimizer.pt`` *inside* a directory. These tests pin the -matching rules, since a wrong match silently builds a preconditioner from the -wrong step rather than failing. -""" - -import os -import shutil -from pathlib import Path - -import pytest -import torch - -from bergson.utils.load_from_optimizer import load_optimizer -from bergson.utils.optimizer_placement import ( - OPTIMIZER_STATE_FILE, - place_final_optimizer_state, - place_optimizer_states, - sorted_optimizer_states, -) - - -def _write_state(path: Path, step: int) -> None: - """A minimal optimizer.pt whose payload identifies its step.""" - path.parent.mkdir(parents=True, exist_ok=True) - torch.save( - { - "state": {0: {"exp_avg_sq": torch.full((2, 3), float(step))}}, - "param_groups": [{"betas": (0.9, 0.999), "eps": 1e-8, "lr": 1e-4}], - }, - path, - ) - - -@pytest.fixture -def save_dir(tmp_path) -> Path: - """A trainer output dir with sibling states at steps 0, 4 and 9.""" - d = tmp_path / "run" - for step in (0, 4, 9): - _write_state(d / f"step_{step}.optimizer.pt", step) - (d / f"step_{step}.ckpt").mkdir(parents=True, exist_ok=True) - return d - - -def _exported(root: Path, names) -> list[Path]: - dirs = [] - for name in names: - p = root / name - p.mkdir(parents=True, exist_ok=True) - dirs.append(p) - return dirs - - -def test_sorted_optimizer_states_orders_numerically(save_dir): - """Steps must sort numerically, not lexically (step_10 < step_9 as strings).""" - _write_state(save_dir / "step_10.optimizer.pt", 10) - assert [s for s, _ in sorted_optimizer_states(save_dir)] == [0, 4, 9, 10] - - -def test_sorted_optimizer_states_ignores_other_files(save_dir): - """The final optimizer.pt and the .ckpt dirs are not per-step states.""" - _write_state(save_dir / OPTIMIZER_STATE_FILE, -1) - assert [s for s, _ in sorted_optimizer_states(save_dir)] == [0, 4, 9] - - -def test_places_by_step_name(save_dir, tmp_path): - """Directories naming a step get that step's state, regardless of order.""" - dsts = _exported(tmp_path / "hf", ["step_9", "step_0", "step_4"]) - placed = place_optimizer_states(save_dir, dsts, mode="copy") - - assert len(placed) == 3 - for dst, expected in zip(dsts, (9, 0, 4)): - blob = load_optimizer(str(dst)) - assert blob["state"][0]["exp_avg_sq"][0, 0].item() == expected - - -def test_places_positionally_when_unnamed(save_dir, tmp_path): - """Unnamed dirs are matched in step order, so seg 0 gets the earliest step.""" - dsts = _exported(tmp_path / "hf", ["a", "b", "c"]) - place_optimizer_states(save_dir, dsts, mode="copy") - - for dst, expected in zip(dsts, (0, 4, 9)): - blob = load_optimizer(str(dst)) - assert blob["state"][0]["exp_avg_sq"][0, 0].item() == expected - - -def test_positional_requires_equal_counts(save_dir, tmp_path): - """A short/long dir list means the caller passed the wrong run; don't guess.""" - dsts = _exported(tmp_path / "hf", ["a", "b"]) - with pytest.raises(ValueError, match="Cannot match positionally"): - place_optimizer_states(save_dir, dsts, mode="copy") - - -def test_mixed_naming_is_rejected(save_dir, tmp_path): - """Half-named lists can't be matched reliably, so refuse rather than guess.""" - dsts = _exported(tmp_path / "hf", ["step_0", "b", "c"]) - with pytest.raises(ValueError, match="name a step and some do not"): - place_optimizer_states(save_dir, dsts, mode="copy") - - -def test_unknown_step_is_rejected(save_dir, tmp_path): - """A named step with no matching state is an error, not a skip.""" - dsts = _exported(tmp_path / "hf", ["step_0", "step_7"]) - with pytest.raises(FileNotFoundError, match=r"steps \[7\]"): - place_optimizer_states(save_dir, dsts, mode="copy") - - -def test_refuses_to_clobber_without_overwrite(save_dir, tmp_path): - """Existing optimizer.pt files are protected; nothing is written on refusal.""" - dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) - (dsts[1] / OPTIMIZER_STATE_FILE).write_text("do not clobber") - - with pytest.raises(FileExistsError): - place_optimizer_states(save_dir, dsts, mode="copy") - - assert (dsts[1] / OPTIMIZER_STATE_FILE).read_text() == "do not clobber" - # The check precedes all writes, so the untouched dirs stay untouched. - assert not (dsts[0] / OPTIMIZER_STATE_FILE).exists() - - place_optimizer_states(save_dir, dsts, mode="copy", overwrite=True) - assert load_optimizer(str(dsts[1]))["state"][0]["exp_avg_sq"][0, 0].item() == 4 - - -def test_symlink_mode_is_relative_and_loadable(save_dir, tmp_path): - """Default mode links rather than duplicating a state per checkpoint.""" - dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) - place_optimizer_states(save_dir, dsts) - - link = dsts[0] / OPTIMIZER_STATE_FILE - assert link.is_symlink() - assert not os.path.isabs(os.readlink(link)), "link must be relative" - assert load_optimizer(str(dsts[0]))["state"][0]["exp_avg_sq"][0, 0].item() == 0 - - -def test_symlinks_survive_relocation(save_dir, tmp_path): - """Relative links keep resolving when run and checkpoints move together.""" - dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) - place_optimizer_states(save_dir, dsts) - - moved = tmp_path / "moved" - moved.mkdir() - shutil.copytree(tmp_path / "hf", moved / "hf", symlinks=True) - shutil.copytree(save_dir, moved / save_dir.name, symlinks=True) - - relocated = moved / "hf" / "step_0" / OPTIMIZER_STATE_FILE - assert relocated.is_symlink() - assert relocated.resolve().is_file(), "relative link broke after relocation" - assert ( - load_optimizer(str(moved / "hf" / "step_0"))["state"][0]["exp_avg_sq"][ - 0, 0 - ].item() - == 0 - ) - - -def test_hardlink_and_move_modes(save_dir, tmp_path): - dsts = _exported(tmp_path / "hf", ["step_0", "step_4", "step_9"]) - place_optimizer_states(save_dir, dsts, mode="hardlink") - assert load_optimizer(str(dsts[2]))["state"][0]["exp_avg_sq"][0, 0].item() == 9 - - dsts2 = _exported(tmp_path / "hf2", ["step_0", "step_4", "step_9"]) - place_optimizer_states(save_dir, dsts2, mode="move") - assert not (save_dir / "step_9.optimizer.pt").exists() - assert load_optimizer(str(dsts2[2]))["state"][0]["exp_avg_sq"][0, 0].item() == 9 - - -def test_missing_states_names_the_config_flag(tmp_path): - """The error should say how to produce the files, not just that they're absent.""" - empty = tmp_path / "empty" - empty.mkdir() - dsts = _exported(tmp_path / "hf", ["step_0"]) - with pytest.raises(FileNotFoundError, match="save_optimizer_state"): - place_optimizer_states(empty, dsts) - - -def test_missing_destination_is_rejected(save_dir, tmp_path): - with pytest.raises(NotADirectoryError): - place_optimizer_states(save_dir, [tmp_path / "nope"]) - - -def test_place_final_state_for_trackstar(save_dir, tmp_path): - """TrackStar points optimizer_state at a dir; the final state must land there.""" - _write_state(save_dir / OPTIMIZER_STATE_FILE, 99) - exported = _exported(tmp_path / "model", ["final"])[0] - - target = place_final_optimizer_state(save_dir, exported, mode="copy") - - assert target == exported / OPTIMIZER_STATE_FILE - # load_optimizer resolves a directory, which is what PreprocessConfig takes. - assert load_optimizer(str(exported))["state"][0]["exp_avg_sq"][0, 0].item() == 99 - - -def test_place_final_state_is_idempotent_in_place(save_dir): - """Pointing it at its own directory is a no-op, not a self-clobber.""" - _write_state(save_dir / OPTIMIZER_STATE_FILE, 99) - target = place_final_optimizer_state(save_dir, save_dir) - assert target.is_file() - assert load_optimizer(str(save_dir))["state"][0]["exp_avg_sq"][0, 0].item() == 99 - - -def test_place_final_state_missing_names_the_flag(tmp_path): - empty = tmp_path / "empty" - empty.mkdir() - dst = _exported(tmp_path / "model", ["final"])[0] - with pytest.raises(FileNotFoundError, match="save_optimizer_state"): - place_final_optimizer_state(empty, dst) diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index 8c75e49b..a76ec7d3 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -308,3 +308,62 @@ def test_lr_history_read_from_the_run_not_the_export(tmp_path): pytest.approx(2e-3), pytest.approx(2e-3), ] + + +def test_dcp_tolerates_optimizer_state_inside_the_checkpoint(tmp_path): + """The layout rests on this: an optimizer.pt inside step_.ckpt/ must not + disturb DCP's own load or a resumed run, and must survive a re-save.""" + import torchopt + from datasets import Dataset + from transformers import AutoConfig, AutoModelForCausalLM + + from bergson.magic.data_stream import DataStream + from bergson.magic.trainer import Trainer + from bergson.utils.trainer_export import OPTIMIZER_STATE_FILE + + config = AutoConfig.from_pretrained("EleutherAI/pythia-14m") + + def fresh(): + torch.manual_seed(0) + m = AutoModelForCausalLM.from_config( + config, dtype=torch.float32, attn_implementation="eager" + ) + m.requires_grad_(True) + return m + + n = 4 + ds = Dataset.from_dict( + {"input_ids": [[1, 2, 3, 4]] * n, "labels": [[1, 2, 3, 4]] * n} + ) + stream = DataStream(ds, batch_size=1, device="cpu") + opt = torchopt.sgd(lambda step: 1e-4, momentum=0.95) + + save_dir = tmp_path / "checkpoints" + trainer, state = Trainer.initialize(fresh(), opt) + final = trainer.train( + state, stream, inplace=True, save_dir=str(save_dir), save_mode="all" + ) + + ckpt = save_dir / "step_2.ckpt" + torch.save( + {"state": {0: {"exp_avg_sq": torch.ones(2, 2)}}}, ckpt / OPTIMIZER_STATE_FILE + ) + + model = fresh() + _, loaded = Trainer.initialize(model, opt) + loaded.load(str(ckpt)) + + trainer2, state2 = Trainer.initialize(fresh(), opt) + resumed = trainer2.train( + state2, + stream, + inplace=True, + save_dir=str(save_dir), + save_mode="all", + resume=True, + ) + for k in final.params: + torch.testing.assert_close(resumed.params[k], final.params[k]) + + blob = torch.load(ckpt / OPTIMIZER_STATE_FILE, weights_only=False) + assert blob["state"][0]["exp_avg_sq"].shape == (2, 2) From 21ab6f8cb34e45a06153bdc47553977615d8715c Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 13:29:42 +0000 Subject: [PATCH 06/13] feat: write each checkpoint's optimizer state inside that checkpoint save_optimizer_state only ever wrote the end-of-training state, so SOURCE's Adam variant had nothing per checkpoint to read. Trainer.train now takes optimizer_cfg and writes each snapshot's second moments to optimizer.pt inside that snapshot's own directory, tagged with the step and the optimizer's betas/eps/eps_root so the file is self-describing. Inside rather than beside: the state belongs to one checkpoint, so keeping it there means nothing downstream has to match a loose file back to a step. The write happens before the DCP save starts, into the directory DCP is about to populate -- safe because DCP leaves unrelated files in a checkpoint directory alone, which is pinned by test_dcp_tolerates_optimizer_state_inside_the_checkpoint. The Adam SOURCE branch still writes the old sibling layout and should be fixed up to match. --- bergson/config/config.py | 7 ++-- bergson/magic/cli.py | 9 +++++ bergson/magic/trainer.py | 23 +++++++++++ tests/test_source_trainer_integration.py | 49 ++++++++++++++++++++++++ 4 files changed, 85 insertions(+), 3 deletions(-) diff --git a/bergson/config/config.py b/bergson/config/config.py index 7cdf43f3..3f210d04 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -353,9 +353,10 @@ class TrainingConfig(AttributionConfig, Serializable): """Enable ``.train()`` mode. Not certified MAGIC-safe.""" save_optimizer_state: bool = False - """After training, export the optimizer's second moments to - ``/optimizer.pt`` for use as an attribution normalizer. - AdamW only.""" + """Persist the optimizer's second moments: the final state at + ``/optimizer.pt`` (an attribution normalizer) and each + checkpoint's state at ``optimizer.pt`` inside that checkpoint's own + directory. AdamW only.""" weight_decay: float = 0.01 """Weight decay coefficient for AdamW and Muon.""" diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index 319fd5ab..3ce10452 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -256,6 +256,15 @@ def worker( fsdp=run_cfg.fsdp, max_grad_norm=run_cfg.max_grad_norm, grad_accum_steps=run_cfg.grad_accum_steps, + optimizer_cfg=( + dict( + betas=(run_cfg.adam_beta1, run_cfg.adam_beta2), + eps=run_cfg.adam_eps, + eps_root=run_cfg.eps_root, + ) + if getattr(run_cfg, "save_optimizer_state", False) + else None + ), ) if getattr(run_cfg, "save_optimizer_state", False) and global_rank == 0: save_second_moments_as_optimizer_pt( diff --git a/bergson/magic/trainer.py b/bergson/magic/trainer.py index 3ce3681e..2c4680ee 100644 --- a/bergson/magic/trainer.py +++ b/bergson/magic/trainer.py @@ -22,6 +22,7 @@ from ..config.config import TrainingConfig from ..data import sorted_checkpoints +from ..utils.load_from_optimizer import save_second_moments_as_optimizer_pt from ..utils.utils import get_device from ..utils.worker_utils import setup_model_and_peft from .config import MagicSaveMode @@ -553,6 +554,7 @@ def train( fsdp: bool = False, max_grad_norm: float | None = None, grad_accum_steps: int = 1, + optimizer_cfg: dict | None = None, ) -> TrainerState: """Train the model on the given data stream, starting from the given state. @@ -578,6 +580,10 @@ def train( step. Passed through to `Trainer.step`. grad_accum_steps: Number of micro-batches to accumulate gradients over per optimizer step. Passed through to `Trainer.step`. + optimizer_cfg: When set (to the optimizer's betas/eps/eps_root), + write each checkpoint's second moments to ``optimizer.pt`` + *inside* that checkpoint's directory, tagged with the step and + these hyperparameters. AdamW only. Returns: The final trainer state after training. @@ -615,6 +621,23 @@ def train( pending_save.result() p = os.path.join(save_dir, f"step_{i}.ckpt") + + # Written before the DCP save starts, into the directory DCP is + # about to populate: the state belongs to this checkpoint, so it + # lives inside it rather than beside it, and consumers never have + # to match a loose file back to a step. DCP leaves unrelated + # files in the directory alone (see test_source_trainer_ + # integration.test_dcp_tolerates_optimizer_state_inside_the_checkpoint). + 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=i, + **optimizer_cfg, + ) + pending_save = state.save(p, debug_pbar=pbar if debug else None) next_save = next_save_index(next_save, n, save_mode) diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index a76ec7d3..dc534b36 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -367,3 +367,52 @@ def fresh(): blob = torch.load(ckpt / OPTIMIZER_STATE_FILE, weights_only=False) assert blob["state"][0]["exp_avg_sq"].shape == (2, 2) + + +def test_trainer_writes_optimizer_state_inside_each_checkpoint(tmp_path): + """optimizer_cfg puts each step's second moments in that step's own dir.""" + import torchopt + from datasets import Dataset + from transformers import AutoConfig, AutoModelForCausalLM + + from bergson.magic.data_stream import DataStream + from bergson.magic.trainer import Trainer + from bergson.utils.load_from_optimizer import load_optimizer + from bergson.utils.trainer_export import ( + sorted_dcp_checkpoints, + ) + + torch.manual_seed(0) + config = AutoConfig.from_pretrained("EleutherAI/pythia-14m") + model = AutoModelForCausalLM.from_config( + config, dtype=torch.float32, attn_implementation="eager" + ) + model.requires_grad_(True) + + n = 3 + ds = Dataset.from_dict( + {"input_ids": [[1, 2, 3, 4]] * n, "labels": [[1, 2, 3, 4]] * n} + ) + stream = DataStream(ds, batch_size=1, device="cpu") + trainer, state = Trainer.initialize(model, torchopt.adamw(lambda step: 1e-4)) + + save_dir = tmp_path / "checkpoints" + trainer.train( + state, + stream, + inplace=True, + save_dir=str(save_dir), + save_mode="all", + optimizer_cfg=dict(betas=(0.9, 0.999), eps=1e-8, eps_root=0.0), + ) + + # No loose siblings: the state lives inside the checkpoint it belongs to. + assert not list(save_dir.glob("step_*.optimizer.pt")) + + for step, ckpt in sorted_dcp_checkpoints(save_dir): + blob = load_optimizer(str(ckpt)) + assert blob["state"], f"step {step} has no second moments" + entry = next(iter(blob["state"].values())) + recorded = entry["step"] + assert int(recorded.item() if torch.is_tensor(recorded) else recorded) == step + assert blob["param_groups"][0]["betas"] == (0.9, 0.999) From ad06dc33390a786f192b85a66f3fa643ccc76adb Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 13:32:36 +0000 Subject: [PATCH 07/13] fix: discovery now finds the export's default destination export_checkpoints defaults to /exported/ but discover_checkpoints only globbed /checkpoint-*, so a default export left resolve() raising 'export them first' against checkpoints that were already there. Both sides now share EXPORT_DIRNAME, and discovery still accepts an export pointed at the run root. --- bergson/approx_unrolling/trainer_run.py | 19 +++++++++++++------ bergson/utils/trainer_export.py | 4 ++-- tests/test_source_trainer_integration.py | 14 ++++++++++++++ 3 files changed, 29 insertions(+), 8 deletions(-) diff --git a/bergson/approx_unrolling/trainer_run.py b/bergson/approx_unrolling/trainer_run.py index 9906b03e..250aba49 100644 --- a/bergson/approx_unrolling/trainer_run.py +++ b/bergson/approx_unrolling/trainer_run.py @@ -15,6 +15,10 @@ from ..config.config_io import CONFIG_FILENAME from ..utils.logger import get_logger +EXPORT_DIRNAME = "exported" +"""Where export_checkpoints puts ``checkpoint-`` dirs by default, and so the +first place discovery looks.""" + LR_HISTORY_FILENAME = "log_history.json" """Per-step LRs in HF's ``log_history`` shape, written beside a run's checkpoints -- the path the LR math already checks first.""" @@ -90,12 +94,15 @@ def discover_checkpoints(trainer_run: str | Path) -> list[str]: """Saved checkpoints in training order. Prefers exported ``checkpoint-`` dirs; finding only native ``step_.ckpt`` raises, naming the export.""" root = Path(trainer_run) - exported = sorted( - (p for p in root.glob("checkpoint-*") if p.is_dir()), - key=lambda p: int(p.name.split("-")[1]), - ) - if exported: - return [str(p) for p in exported] + # The export's default destination first, then the run root for an export + # that was pointed there explicitly. + for base in (root / EXPORT_DIRNAME, root): + exported = sorted( + (p for p in base.glob("checkpoint-*") if p.is_dir()), + key=lambda p: int(p.name.split("-")[1]), + ) + if exported: + return [str(p) for p in exported] ckpt_dir = root / "checkpoints" native = sorted( diff --git a/bergson/utils/trainer_export.py b/bergson/utils/trainer_export.py index f4140eab..78011f58 100644 --- a/bergson/utils/trainer_export.py +++ b/bergson/utils/trainer_export.py @@ -7,7 +7,7 @@ import torch from transformers import AutoTokenizer -from ..approx_unrolling.trainer_run import load_training_config +from ..approx_unrolling.trainer_run import EXPORT_DIRNAME, load_training_config from ..config.config import TrainingConfig from ..magic.trainer import prepare_trainer from .logger import get_logger @@ -50,7 +50,7 @@ def export_checkpoints( export, since each one is a full model copy. """ run_path = Path(run_path) - out_dir = Path(out_dir) if out_dir is not None else run_path / "exported" + out_dir = Path(out_dir) if out_dir is not None else run_path / EXPORT_DIRNAME cfg = training_cfg if training_cfg is not None else load_training_config(run_path) available = sorted_dcp_checkpoints(run_path / "checkpoints") diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index dc534b36..b2043adb 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -416,3 +416,17 @@ def test_trainer_writes_optimizer_state_inside_each_checkpoint(tmp_path): recorded = entry["step"] assert int(recorded.item() if torch.is_tensor(recorded) else recorded) == step assert blob["param_groups"][0]["betas"] == (0.9, 0.999) + + +def test_discovery_finds_the_default_export_destination(tmp_path): + """resolve() must find what export_checkpoints writes with no out_dir given.""" + from bergson.approx_unrolling.trainer_run import EXPORT_DIRNAME + + run = _run_dir(tmp_path) + export = run / EXPORT_DIRNAME + for step in (0, 10, 2): + (export / f"checkpoint-{step}").mkdir(parents=True) + + out = resolve(ApproxUnrollingConfig(trainer_run=str(run))) + + assert [p.split("-")[-1] for p in out.checkpoints] == ["0", "2", "10"] From a7402a7c8e4197635cbb2f894d74fd78abb81709 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 13:42:32 +0000 Subject: [PATCH 08/13] refactor: drop unused _checkpoint_dir_step Never called, and duplicated _checkpoint_step in approx_unrolling_math. --- bergson/approx_unrolling/trainer_run.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/bergson/approx_unrolling/trainer_run.py b/bergson/approx_unrolling/trainer_run.py index 250aba49..ca4953fe 100644 --- a/bergson/approx_unrolling/trainer_run.py +++ b/bergson/approx_unrolling/trainer_run.py @@ -5,7 +5,6 @@ """ import json -import os from pathlib import Path from typing import Any, Callable @@ -167,16 +166,3 @@ def lr_history_path(cfg: ApproxUnrollingConfig) -> Path | None: if candidate.is_file(): return candidate return None - - -def _checkpoint_dir_step(path: str | os.PathLike) -> int | None: - """Step index named by a checkpoint dir, across the layouts we emit/accept.""" - name = Path(path).name - if name.startswith("checkpoint-") and name.removeprefix("checkpoint-").isdigit(): - return int(name.removeprefix("checkpoint-")) - stem = name.removesuffix(".ckpt") - if stem.startswith("step_") and stem.removeprefix("step_").isdigit(): - return int(stem.removeprefix("step_")) - if name.isdigit(): - return int(name) - return None From 4e033b7cc5b47d30151864a9fba7b43b7f26185f Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 13:46:01 +0000 Subject: [PATCH 09/13] feat: infer trainer_run from the checkpoints' path Pointing at a bergson run twice was redundant: an exported checkpoint already says which run it came from. Setting checkpoints is now enough; trainer_run is only needed to discover them or to name a different run. Only the two layouts we emit are probed (/exported/checkpoint-N and /checkpoint-N), so an unrelated config.yaml higher up is never mistaken for the producing run, and foreign checkpoints infer nothing. --- bergson/approx_unrolling/trainer_run.py | 25 +++++++++++++++-- bergson/config/config.py | 5 ++-- tests/test_source_trainer_integration.py | 35 ++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 5 deletions(-) diff --git a/bergson/approx_unrolling/trainer_run.py b/bergson/approx_unrolling/trainer_run.py index ca4953fe..59d215cd 100644 --- a/bergson/approx_unrolling/trainer_run.py +++ b/bergson/approx_unrolling/trainer_run.py @@ -123,16 +123,35 @@ def discover_checkpoints(trainer_run: str | Path) -> list[str]: ) +def infer_trainer_run(checkpoints: list[str]) -> str: + """The bergson run a checkpoint came from, or "" if it did not come from one. + + Only the two layouts we emit are considered -- ``/exported/checkpoint-N`` + and ``/checkpoint-N`` -- so an unrelated config.yaml further up the tree + is never mistaken for the run that produced these checkpoints. Checkpoints + from other trainers simply find nothing. + """ + if not checkpoints: + return "" + first = Path(checkpoints[0]) + for candidate in (first.parent.parent, first.parent): + if (candidate / CONFIG_FILENAME).is_file(): + return str(candidate) + return "" + + def resolve(cfg: ApproxUnrollingConfig) -> ApproxUnrollingConfig: """Fill unset fields from ``cfg.trainer_run``. A no-op when it is empty; never overwrites a field the caller set.""" - if not cfg.trainer_run: - # Nothing to derive from; only normalize the momentum sentinel. + trainer_run = cfg.trainer_run or infer_trainer_run(cfg.checkpoints) + if not trainer_run: + # Not a bergson run; only normalize the momentum sentinel. if cfg.momentum is None: cfg.momentum = 0.0 return cfg + cfg.trainer_run = trainer_run - training_cfg = load_training_config(cfg.trainer_run) + training_cfg = load_training_config(trainer_run) filled: list[str] = [] if not cfg.checkpoints: diff --git a/bergson/config/config.py b/bergson/config/config.py index 3f210d04..e814027f 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -789,8 +789,9 @@ class ApproxUnrollingConfig(Serializable): trainer_run: str = "" """Bergson run directory to read ``checkpoints``, ``model_path`` and - ``momentum`` from. Explicit fields always win; leave empty for other - trainers.""" + ``momentum`` from. Inferred from ``checkpoints`` when they came from one, + so it is only needed to discover the checkpoints themselves or to point at + a different run. Explicit fields always win.""" momentum: float | None = None """SGD heavy-ball momentum beta; scales lr*steps by 1/(1-beta) (Bae et al. diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index b2043adb..85cfbcb6 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -430,3 +430,38 @@ def test_discovery_finds_the_default_export_destination(tmp_path): out = resolve(ApproxUnrollingConfig(trainer_run=str(run))) assert [p.split("-")[-1] for p in out.checkpoints] == ["0", "2", "10"] + + +def test_trainer_run_inferred_from_exported_checkpoints(tmp_path): + """Setting checkpoints is enough; the run is found from their path.""" + from bergson.approx_unrolling.trainer_run import EXPORT_DIRNAME + + run = _run_dir(tmp_path, optimizer="sgd", adam_beta1=0.9, model="gpt2") + ckpt = run / EXPORT_DIRNAME / "checkpoint-3" + ckpt.mkdir(parents=True) + + out = resolve(ApproxUnrollingConfig(checkpoints=[str(ckpt)])) + + assert out.trainer_run == str(run) + assert out.momentum == 0.9 + assert out.model_path == "gpt2" + + +def test_trainer_run_inferred_from_run_root_checkpoints(tmp_path): + run = _run_dir(tmp_path, optimizer="sgd", adam_beta1=0.8) + ckpt = run / "checkpoint-1" + ckpt.mkdir() + + assert resolve(ApproxUnrollingConfig(checkpoints=[str(ckpt)])).momentum == 0.8 + + +def test_foreign_checkpoints_infer_nothing(tmp_path): + """HF Trainer checkpoints must not pick up an unrelated run's config.""" + hf = tmp_path / "hf_run" / "checkpoint-500" + hf.mkdir(parents=True) + + out = resolve(ApproxUnrollingConfig(checkpoints=[str(hf)])) + + assert out.trainer_run == "" + assert out.momentum == 0.0 + assert out.model_path is None From 12fb36766d75a8995c291c35b8bb14793df3ae6b Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 13:53:53 +0000 Subject: [PATCH 10/13] feat: choose which checkpoints get optimizer states save_optimizer_state was a bool that, after the per-checkpoint write landed, meant 'all checkpoints' -- no way to ask for just the final state, and no way to opt out of one file per checkpoint when only the normalizer was wanted. Now none/last/all. 'last' is the original meaning (final state at /optimizer.pt, the attribution normalizer); 'all' adds each checkpoint's own optimizer.pt, which SOURCE's Adam variant needs and which costs a file the size of the trainable parameters per checkpoint. Old booleans still parse: true -> 'last', false -> 'none', so existing configs keep their original behaviour rather than silently gaining per-checkpoint files. --- bergson/config/config.py | 24 +++++++++++++++++++----- bergson/magic/cli.py | 4 ++-- tests/test_source_trainer_integration.py | 10 ++++++++++ 3 files changed, 31 insertions(+), 7 deletions(-) diff --git a/bergson/config/config.py b/bergson/config/config.py index e814027f..e5218b60 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -352,11 +352,17 @@ class TrainingConfig(AttributionConfig, Serializable): train_mode: bool = False """Enable ``.train()`` mode. Not certified MAGIC-safe.""" - save_optimizer_state: bool = False - """Persist the optimizer's second moments: the final state at - ``/optimizer.pt`` (an attribution normalizer) and each - checkpoint's state at ``optimizer.pt`` inside that checkpoint's own - directory. AdamW only.""" + save_optimizer_state: Literal["none", "last", "all"] = "none" + """Which checkpoints' optimizer second moments to persist. AdamW only. + + - "none" (default): none. + - "last": the final state at ``/optimizer.pt``, for use as an + attribution normalizer. + - "all": that, plus each checkpoint's state at ``optimizer.pt`` inside its + own directory -- what SOURCE's Adam variant reads. One extra file per + checkpoint, each the size of the trainable parameters. + + Accepts the old booleans: ``true`` means "last", ``false`` means "none".""" weight_decay: float = 0.01 """Weight decay coefficient for AdamW and Muon.""" @@ -382,6 +388,14 @@ class TrainingConfig(AttributionConfig, Serializable): wandb_project: str = "" """Weights & Biases project name. If set, logs training loss to W&B.""" + def __post_init__(self): + super().__post_init__() + + # Back-compat: save_optimizer_state used to be a bool meaning + # "write the final state". + if isinstance(self.save_optimizer_state, bool): + self.save_optimizer_state = "last" if self.save_optimizer_state else "none" + @dataclass class MetasmoothnessConfig(TrainingConfig): diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index 3ce10452..4dc6b793 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -262,11 +262,11 @@ def worker( eps=run_cfg.adam_eps, eps_root=run_cfg.eps_root, ) - if getattr(run_cfg, "save_optimizer_state", False) + if getattr(run_cfg, "save_optimizer_state", "none") == "all" else None ), ) - if getattr(run_cfg, "save_optimizer_state", False) and global_rank == 0: + if getattr(run_cfg, "save_optimizer_state", "none") != "none" and global_rank == 0: save_second_moments_as_optimizer_pt( model, # type: ignore[reportArgumentType] fwd_state.opt_state, diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index 85cfbcb6..f75f7775 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -465,3 +465,13 @@ def test_foreign_checkpoints_infer_nothing(tmp_path): assert out.trainer_run == "" assert out.momentum == 0.0 assert out.model_path is None + + +@pytest.mark.parametrize( + "value,expected", + [(True, "last"), (False, "none"), ("all", "all"), ("last", "last")], +) +def test_save_optimizer_state_accepts_old_booleans(value, expected): + """`true` used to mean "write the final state"; keep that meaning.""" + cfg = TrainingConfig(run_path="/tmp/x", save_optimizer_state=value) + assert cfg.save_optimizer_state == expected From 7ca4546e8e388358bcf45d37ebcad970a51893e6 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 13:57:19 +0000 Subject: [PATCH 11/13] refactor: drop trainer_run; derive the run from the checkpoints You always say which checkpoints to attribute, and an exported checkpoint already records which run produced it, so naming the run again was redundant. The field and its discovery path are gone; resolve() infers the run from checkpoints[0] and fills only model_path and momentum. Raw step_.ckpt paths now raise up front naming export_checkpoints, instead of failing later inside from_pretrained -- the one useful thing the discovery path did. --- .../approx_unrolling/approx_unrolling_math.py | 4 +- bergson/approx_unrolling/trainer_run.py | 62 +++++----------- bergson/config/config.py | 9 +-- tests/test_source_trainer_integration.py | 71 ++++++------------- 4 files changed, 43 insertions(+), 103 deletions(-) diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index 411bdc31..bf093272 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -79,9 +79,9 @@ def compute_lr_times_steps_per_segment( per_segment = len(cfg.checkpoints) // L ckpt_steps = [_checkpoint_step(p) for p in cfg.checkpoints] boundaries = [0] + [ckpt_steps[(l + 1) * per_segment - 1] for l in range(L)] - # A bergson run's own history, else one dumped beside the checkpoints, + # The bergson run these came from, else a history dumped beside them, # else the final checkpoint's trainer_state.json (what HF Trainer writes). - log_path = lr_history_path(cfg) or ( + log_path = lr_history_path(cfg.checkpoints) or ( Path(str(cfg.checkpoints[0])).parent / LR_HISTORY_FILENAME ) if log_path.exists(): diff --git a/bergson/approx_unrolling/trainer_run.py b/bergson/approx_unrolling/trainer_run.py index 59d215cd..53679e95 100644 --- a/bergson/approx_unrolling/trainer_run.py +++ b/bergson/approx_unrolling/trainer_run.py @@ -89,39 +89,17 @@ def derive_momentum(training_cfg: TrainingConfig) -> float: return 0.0 -def discover_checkpoints(trainer_run: str | Path) -> list[str]: - """Saved checkpoints in training order. Prefers exported ``checkpoint-`` - dirs; finding only native ``step_.ckpt`` raises, naming the export.""" - root = Path(trainer_run) - # The export's default destination first, then the run root for an export - # that was pointed there explicitly. - for base in (root / EXPORT_DIRNAME, root): - exported = sorted( - (p for p in base.glob("checkpoint-*") if p.is_dir()), - key=lambda p: int(p.name.split("-")[1]), - ) - if exported: - return [str(p) for p in exported] - - ckpt_dir = root / "checkpoints" - native = sorted( - (p for p in ckpt_dir.glob("step_*.ckpt") if p.is_dir()), - key=lambda p: int(p.name.removesuffix(".ckpt").removeprefix("step_")), - ) +def _reject_unexported(checkpoints: list[str]) -> None: + """SOURCE loads checkpoints with from_pretrained, so a raw DCP directory + would fail deep in the pipeline; say what to do instead.""" + native = [c for c in checkpoints if Path(c).name.endswith(".ckpt")] if native: - raise FileNotFoundError( - f"{trainer_run} has {len(native)} native checkpoints under " - f"{ckpt_dir} but no exported checkpoint- directories. SOURCE " - "loads models with from_pretrained, so export them first with " - "bergson.utils.trainer_export.export_checkpoints(run_path, out_dir)." + raise ValueError( + f"{native[:3]} are the trainer's DCP checkpoints, which " + "from_pretrained cannot load. Export them first with " + "bergson.utils.trainer_export.export_checkpoints(run_path)." ) - raise FileNotFoundError( - f"No checkpoints found under {trainer_run}. Train with a save_mode " - "that writes them, then export with " - "bergson.utils.trainer_export.export_checkpoints." - ) - def infer_trainer_run(checkpoints: list[str]) -> str: """The bergson run a checkpoint came from, or "" if it did not come from one. @@ -143,21 +121,18 @@ def infer_trainer_run(checkpoints: list[str]) -> str: def resolve(cfg: ApproxUnrollingConfig) -> ApproxUnrollingConfig: """Fill unset fields from ``cfg.trainer_run``. A no-op when it is empty; never overwrites a field the caller set.""" - trainer_run = cfg.trainer_run or infer_trainer_run(cfg.checkpoints) + _reject_unexported(cfg.checkpoints) + + trainer_run = infer_trainer_run(cfg.checkpoints) if not trainer_run: # Not a bergson run; only normalize the momentum sentinel. if cfg.momentum is None: cfg.momentum = 0.0 return cfg - cfg.trainer_run = trainer_run training_cfg = load_training_config(trainer_run) filled: list[str] = [] - if not cfg.checkpoints: - cfg.checkpoints = discover_checkpoints(cfg.trainer_run) - filled.append(f"checkpoints ({len(cfg.checkpoints)})") - if cfg.model_path is None: cfg.model_path = training_cfg.model filled.append(f"model_path={cfg.model_path!r}") @@ -167,20 +142,19 @@ def resolve(cfg: ApproxUnrollingConfig) -> ApproxUnrollingConfig: filled.append(f"momentum={cfg.momentum}") if filled: - logger.info( - "Filled from trainer_run %s: %s", cfg.trainer_run, ", ".join(filled) - ) + logger.info("Filled from run %s: %s", trainer_run, ", ".join(filled)) return cfg -def lr_history_path(cfg: ApproxUnrollingConfig) -> Path | None: - """Where a bergson run's LR history lives, if this config names one.""" - if not cfg.trainer_run: +def lr_history_path(checkpoints: list[str]) -> Path | None: + """The LR history of the bergson run these checkpoints came from, if any.""" + run = infer_trainer_run(checkpoints) + if not run: return None for candidate in ( - Path(cfg.trainer_run) / LR_HISTORY_FILENAME, - Path(cfg.trainer_run) / "checkpoints" / LR_HISTORY_FILENAME, + Path(run) / LR_HISTORY_FILENAME, + Path(run) / "checkpoints" / LR_HISTORY_FILENAME, ): if candidate.is_file(): return candidate diff --git a/bergson/config/config.py b/bergson/config/config.py index e5218b60..4d4867d5 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -801,15 +801,10 @@ class ApproxUnrollingConfig(Serializable): """Per-segment optimizer step count; length must == segments. If empty, inferred from per-checkpoint step counts.""" - trainer_run: str = "" - """Bergson run directory to read ``checkpoints``, ``model_path`` and - ``momentum`` from. Inferred from ``checkpoints`` when they came from one, - so it is only needed to discover the checkpoints themselves or to point at - a different run. Explicit fields always win.""" - momentum: float | None = None """SGD heavy-ball momentum beta; scales lr*steps by 1/(1-beta) (Bae et al. - 2024, App. D.2). ``None`` derives it from ``trainer_run``, else 0.0. + 2024, App. D.2). ``None`` derives it from the run the checkpoints + came from, else 0.0. bergson's SGD uses ``adam_beta1`` (default 0.95), so assuming 0.0 there understates lr*steps by 20x.""" diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index f75f7775..73482215 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -96,7 +96,9 @@ def test_resolve_fills_momentum_and_model_from_run(tmp_path): run = _run_dir( tmp_path, optimizer="sgd", adam_beta1=0.9, model="EleutherAI/pythia-14m" ) - cfg = ApproxUnrollingConfig(trainer_run=str(run), checkpoints=["a", "b"]) + ckpt = run / "exported" / "checkpoint-0" + ckpt.mkdir(parents=True) + cfg = ApproxUnrollingConfig(checkpoints=[str(ckpt)]) out = resolve(cfg) @@ -107,52 +109,29 @@ def test_resolve_fills_momentum_and_model_from_run(tmp_path): def test_explicit_momentum_wins_over_run(tmp_path): """A user training elsewhere must be able to override what the run says.""" run = _run_dir(tmp_path, optimizer="sgd", adam_beta1=0.9) - cfg = ApproxUnrollingConfig(trainer_run=str(run), checkpoints=["a"], momentum=0.5) + ckpt = run / "checkpoint-0" + ckpt.mkdir() + cfg = ApproxUnrollingConfig(checkpoints=[str(ckpt)], momentum=0.5) assert resolve(cfg).momentum == 0.5 def test_explicit_zero_momentum_is_respected(tmp_path): """0.0 is a real value, not 'unset' -- it must survive derivation.""" run = _run_dir(tmp_path, optimizer="sgd", adam_beta1=0.9) - cfg = ApproxUnrollingConfig(trainer_run=str(run), checkpoints=["a"], momentum=0.0) + ckpt = run / "checkpoint-0" + ckpt.mkdir() + cfg = ApproxUnrollingConfig(checkpoints=[str(ckpt)], momentum=0.0) assert resolve(cfg).momentum == 0.0 def test_explicit_model_path_wins(tmp_path): run = _run_dir(tmp_path, model="EleutherAI/pythia-14m") - cfg = ApproxUnrollingConfig( - trainer_run=str(run), checkpoints=["a"], model_path="gpt2" - ) + ckpt = run / "checkpoint-0" + ckpt.mkdir() + cfg = ApproxUnrollingConfig(checkpoints=[str(ckpt)], model_path="gpt2") assert resolve(cfg).model_path == "gpt2" -def test_explicit_checkpoints_win(tmp_path): - run = _run_dir(tmp_path) - (run / "checkpoint-0").mkdir() - cfg = ApproxUnrollingConfig(trainer_run=str(run), checkpoints=["mine"]) - assert resolve(cfg).checkpoints == ["mine"] - - -def test_resolve_discovers_exported_checkpoints_in_order(tmp_path): - run = _run_dir(tmp_path) - for step in (0, 10, 2): - (run / f"checkpoint-{step}").mkdir() - - out = resolve(ApproxUnrollingConfig(trainer_run=str(run))) - - assert [p.split("-")[-1] for p in out.checkpoints] == ["0", "2", "10"] - - -def test_unexported_run_names_the_export_step(tmp_path): - """Native DCP checkpoints can't be loaded by from_pretrained; say so.""" - run = _run_dir(tmp_path) - (run / "checkpoints").mkdir() - (run / "checkpoints" / "step_0.ckpt").mkdir() - - with pytest.raises(FileNotFoundError, match="export_checkpoints"): - resolve(ApproxUnrollingConfig(trainer_run=str(run))) - - def test_missing_config_yaml_is_explained(tmp_path): with pytest.raises(FileNotFoundError, match="bergson run directory"): load_training_config(tmp_path) @@ -299,7 +278,6 @@ def test_lr_history_read_from_the_run_not_the_export(tmp_path): assert not (export / LR_HISTORY_FILENAME).exists() cfg = ApproxUnrollingConfig( - trainer_run=str(run), checkpoints=[str(export / "checkpoint-2"), str(export / "checkpoint-4")], segments=2, momentum=0.0, @@ -418,20 +396,6 @@ def test_trainer_writes_optimizer_state_inside_each_checkpoint(tmp_path): assert blob["param_groups"][0]["betas"] == (0.9, 0.999) -def test_discovery_finds_the_default_export_destination(tmp_path): - """resolve() must find what export_checkpoints writes with no out_dir given.""" - from bergson.approx_unrolling.trainer_run import EXPORT_DIRNAME - - run = _run_dir(tmp_path) - export = run / EXPORT_DIRNAME - for step in (0, 10, 2): - (export / f"checkpoint-{step}").mkdir(parents=True) - - out = resolve(ApproxUnrollingConfig(trainer_run=str(run))) - - assert [p.split("-")[-1] for p in out.checkpoints] == ["0", "2", "10"] - - def test_trainer_run_inferred_from_exported_checkpoints(tmp_path): """Setting checkpoints is enough; the run is found from their path.""" from bergson.approx_unrolling.trainer_run import EXPORT_DIRNAME @@ -442,7 +406,6 @@ def test_trainer_run_inferred_from_exported_checkpoints(tmp_path): out = resolve(ApproxUnrollingConfig(checkpoints=[str(ckpt)])) - assert out.trainer_run == str(run) assert out.momentum == 0.9 assert out.model_path == "gpt2" @@ -462,7 +425,6 @@ def test_foreign_checkpoints_infer_nothing(tmp_path): out = resolve(ApproxUnrollingConfig(checkpoints=[str(hf)])) - assert out.trainer_run == "" assert out.momentum == 0.0 assert out.model_path is None @@ -475,3 +437,12 @@ def test_save_optimizer_state_accepts_old_booleans(value, expected): """`true` used to mean "write the final state"; keep that meaning.""" cfg = TrainingConfig(run_path="/tmp/x", save_optimizer_state=value) assert cfg.save_optimizer_state == expected + + +def test_dcp_checkpoints_are_rejected_with_the_export_hint(tmp_path): + """Passing raw step_.ckpt dirs should fail early, not inside the pipeline.""" + ckpt = tmp_path / "checkpoints" / "step_0.ckpt" + ckpt.mkdir(parents=True) + + with pytest.raises(ValueError, match="export_checkpoints"): + resolve(ApproxUnrollingConfig(checkpoints=[str(ckpt)])) From c2f252e7e81e7ac0d558975dd2b0e49631d19e17 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Wed, 29 Jul 2026 14:04:10 +0000 Subject: [PATCH 12/13] test: cover export_checkpoints itself The round-trip test proved the save/load mechanism but not the function's wiring: reading the run's config.yaml, building the model through prepare_trainer, honouring steps=, and landing in the default export dir with each checkpoint's optimizer state alongside. Ends by feeding the result to resolve(), so the train -> export -> SOURCE path is covered end to end. --- bergson/approx_unrolling/trainer_run.py | 7 +-- bergson/config/config.py | 20 ++------ bergson/magic/cli.py | 3 -- bergson/magic/trainer.py | 12 ++--- bergson/utils/trainer_export.py | 10 ++-- tests/test_source_trainer_integration.py | 58 ++++++++++++++++++++++++ 6 files changed, 72 insertions(+), 38 deletions(-) diff --git a/bergson/approx_unrolling/trainer_run.py b/bergson/approx_unrolling/trainer_run.py index 53679e95..5e37dd9d 100644 --- a/bergson/approx_unrolling/trainer_run.py +++ b/bergson/approx_unrolling/trainer_run.py @@ -1,8 +1,5 @@ -"""Fill in SOURCE hyperparameters from a bergson run's ``config.yaml``. - -SOURCE was built for HF Trainer runs. Everything here is a fallback: explicit -config fields always win, so other trainers are unaffected. -""" +"""Fill in unset SOURCE configuration from a bergson run's ``config.yaml`` +if present.""" import json from pathlib import Path diff --git a/bergson/config/config.py b/bergson/config/config.py index 4d4867d5..b11a117e 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -353,16 +353,7 @@ class TrainingConfig(AttributionConfig, Serializable): """Enable ``.train()`` mode. Not certified MAGIC-safe.""" save_optimizer_state: Literal["none", "last", "all"] = "none" - """Which checkpoints' optimizer second moments to persist. AdamW only. - - - "none" (default): none. - - "last": the final state at ``/optimizer.pt``, for use as an - attribution normalizer. - - "all": that, plus each checkpoint's state at ``optimizer.pt`` inside its - own directory -- what SOURCE's Adam variant reads. One extra file per - checkpoint, each the size of the trainable parameters. - - Accepts the old booleans: ``true`` means "last", ``false`` means "none".""" + """Which checkpoints' optimizer second moments to persist. AdamW only.""" weight_decay: float = 0.01 """Weight decay coefficient for AdamW and Muon.""" @@ -391,8 +382,7 @@ class TrainingConfig(AttributionConfig, Serializable): def __post_init__(self): super().__post_init__() - # Back-compat: save_optimizer_state used to be a bool meaning - # "write the final state". + # TODO Lucia Quirke August: remove bwd-compat if isinstance(self.save_optimizer_state, bool): self.save_optimizer_state = "last" if self.save_optimizer_state else "none" @@ -803,10 +793,8 @@ class ApproxUnrollingConfig(Serializable): momentum: float | None = None """SGD heavy-ball momentum beta; scales lr*steps by 1/(1-beta) (Bae et al. - 2024, App. D.2). ``None`` derives it from the run the checkpoints - came from, else 0.0. - bergson's SGD uses ``adam_beta1`` (default 0.95), so assuming 0.0 there - understates lr*steps by 20x.""" + 2024, App. D.2). When ``None`` it's derived from the run if present or set + to 0.0.""" query: DataConfig = field(default_factory=DataConfig) """Query dataset spec; gradients computed at the final checkpoint.""" diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index 4dc6b793..9f0e5979 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -238,9 +238,6 @@ def worker( ckpts_path = os.path.join(run_cfg.run_path, "checkpoints") resume = run_cfg.resume - # Record the realized LR schedule beside the checkpoints, in HF's - # log_history shape, so SOURCE can read per-step LRs off a bergson run - # through the same path it already uses for HF Trainer runs. if global_rank == 0: write_lr_history(ckpts_path, schedule, len(stream)) diff --git a/bergson/magic/trainer.py b/bergson/magic/trainer.py index 2c4680ee..61748371 100644 --- a/bergson/magic/trainer.py +++ b/bergson/magic/trainer.py @@ -581,9 +581,8 @@ def train( grad_accum_steps: Number of micro-batches to accumulate gradients over per optimizer step. Passed through to `Trainer.step`. optimizer_cfg: When set (to the optimizer's betas/eps/eps_root), - write each checkpoint's second moments to ``optimizer.pt`` - *inside* that checkpoint's directory, tagged with the step and - these hyperparameters. AdamW only. + write each checkpoint's second moments to ``optimizer.pt``, + tagged with the step and these hyperparameters. AdamW only. Returns: The final trainer state after training. @@ -622,12 +621,7 @@ def train( p = os.path.join(save_dir, f"step_{i}.ckpt") - # Written before the DCP save starts, into the directory DCP is - # about to populate: the state belongs to this checkpoint, so it - # lives inside it rather than beside it, and consumers never have - # to match a loose file back to a step. DCP leaves unrelated - # files in the directory alone (see test_source_trainer_ - # integration.test_dcp_tolerates_optimizer_state_inside_the_checkpoint). + # Write before the DCP save starts, into the same directory. if optimizer_cfg is not None: os.makedirs(p, exist_ok=True) save_second_moments_as_optimizer_pt( diff --git a/bergson/utils/trainer_export.py b/bergson/utils/trainer_export.py index 78011f58..93b5d932 100644 --- a/bergson/utils/trainer_export.py +++ b/bergson/utils/trainer_export.py @@ -44,10 +44,10 @@ def export_checkpoints( steps: list[int] | None = None, overwrite: bool = False, ) -> list[Path]: - """Export ``step_.ckpt`` to ``checkpoint-/`` dirs, in step order. + """Export ``step_.ckpt`` to ``checkpoint-/`` HF dirs, in step order. - ``training_cfg`` defaults to the run's ``config.yaml``. ``steps`` limits the - export, since each one is a full model copy. + ``training_cfg`` defaults to the run's ``config.yaml``. ``steps`` limits which + checkpoints are exported. """ run_path = Path(run_path) out_dir = Path(out_dir) if out_dir is not None else run_path / EXPORT_DIRNAME @@ -80,8 +80,8 @@ def export_checkpoints( f"{[str(c) for c in clashes]} already exist; pass overwrite=True." ) - # One model/state is built and reloaded per checkpoint, rather than one per - # step: the state's tensors are loaded in place by TrainerState.load. + # One model and state is built and reloaded per checkpoint. + # The state's tensors are loaded in place by TrainerState.load. trainer, state, model = prepare_trainer(cfg, rank=0, schedule=lambda step: 0.0) tokenizer = AutoTokenizer.from_pretrained(cfg.model) diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py index 73482215..f562fef2 100644 --- a/tests/test_source_trainer_integration.py +++ b/tests/test_source_trainer_integration.py @@ -446,3 +446,61 @@ def test_dcp_checkpoints_are_rejected_with_the_export_hint(tmp_path): with pytest.raises(ValueError, match="export_checkpoints"): resolve(ApproxUnrollingConfig(checkpoints=[str(ckpt)])) + + +def test_export_checkpoints_end_to_end(tmp_path): + """The function itself: reads the run's config.yaml, builds the model via + prepare_trainer, and writes loadable checkpoint-/ dirs with their + optimizer state -- the wiring the round-trip test does not cover.""" + import torchopt + from datasets import Dataset + from transformers import AutoConfig, AutoModelForCausalLM + + from bergson.approx_unrolling.trainer_run import EXPORT_DIRNAME + from bergson.magic.data_stream import DataStream + from bergson.magic.trainer import Trainer + from bergson.utils.load_from_optimizer import load_optimizer + from bergson.utils.trainer_export import ( + OPTIMIZER_STATE_FILE, + export_checkpoints, + ) + + model_name = "EleutherAI/pythia-14m" + run = _run_dir(tmp_path, model=model_name, optimizer="adamw") + + torch.manual_seed(0) + config = AutoConfig.from_pretrained(model_name) + model = AutoModelForCausalLM.from_config( + config, dtype=torch.float32, attn_implementation="eager" + ) + model.requires_grad_(True) + + n = 3 + ds = Dataset.from_dict( + {"input_ids": [[1, 2, 3, 4]] * n, "labels": [[1, 2, 3, 4]] * n} + ) + stream = DataStream(ds, batch_size=1, device="cpu") + trainer, state = Trainer.initialize(model, torchopt.adamw(lambda step: 1e-4)) + trainer.train( + state, + stream, + inplace=True, + save_dir=str(run / "checkpoints"), + save_mode="all", + optimizer_cfg=dict(betas=(0.9, 0.999), eps=1e-8, eps_root=0.0), + ) + + exported = export_checkpoints(run, steps=[0, 2]) + + assert [p.name for p in exported] == ["checkpoint-0", "checkpoint-2"] + assert exported[0].parent == run / EXPORT_DIRNAME + + for dst in exported: + AutoModelForCausalLM.from_pretrained(str(dst)) + assert load_optimizer(str(dst))["state"], f"{dst} lost its optimizer state" + assert (dst / OPTIMIZER_STATE_FILE).is_file() + + # And the exported dirs are what resolve() then picks up. + out = resolve(ApproxUnrollingConfig(checkpoints=[str(p) for p in exported])) + assert out.model_path == model_name + assert out.momentum == 0.0 # adamw From 47382c63bbd87d316dd5c2bf17ff4d2bc0ef71ee Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:15:25 +0000 Subject: [PATCH 13/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- bergson/magic/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bergson/magic/trainer.py b/bergson/magic/trainer.py index 61748371..6a92c33a 100644 --- a/bergson/magic/trainer.py +++ b/bergson/magic/trainer.py @@ -581,7 +581,7 @@ def train( grad_accum_steps: Number of micro-batches to accumulate gradients over per optimizer step. Passed through to `Trainer.step`. optimizer_cfg: When set (to the optimizer's betas/eps/eps_root), - write each checkpoint's second moments to ``optimizer.pt``, + write each checkpoint's second moments to ``optimizer.pt``, tagged with the step and these hyperparameters. AdamW only. Returns: