diff --git a/bergson/approx_unrolling/approx_unrolling_math.py b/bergson/approx_unrolling/approx_unrolling_math.py index eab33bc7..bf093272 100644 --- a/bergson/approx_unrolling/approx_unrolling_math.py +++ b/bergson/approx_unrolling/approx_unrolling_math.py @@ -24,24 +24,31 @@ 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. Accepts a path whose final component is a ``checkpoint-`` directory - (what HF Trainer writes natively) 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 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 " "`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." @@ -49,23 +56,34 @@ 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.""" - cfg = approx_unrolling_cfg + 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). + """ 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] 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" + # 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.checkpoints) 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) @@ -75,7 +93,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..a2853c8b 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 # 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,9 @@ def approx_unrolling_pipeline( """ logger = get_logger("approx_unrolling_pipeline") + # 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 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..5e37dd9d --- /dev/null +++ b/bergson/approx_unrolling/trainer_run.py @@ -0,0 +1,158 @@ +"""Fill in unset SOURCE configuration from a bergson run's ``config.yaml`` +if present.""" + +import json +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 + +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.""" + +logger = get_logger(__name__) + + +def write_lr_history( + save_dir: str | Path, schedule: Callable[[int], float], num_steps: int +) -> Path: + """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 = [ + {"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 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( + 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: + """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) + 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 _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 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)." + ) + + +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.""" + _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 + + training_cfg = load_training_config(trainer_run) + filled: list[str] = [] + + 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 run %s: %s", trainer_run, ", ".join(filled)) + + return cfg + + +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(run) / LR_HISTORY_FILENAME, + Path(run) / "checkpoints" / LR_HISTORY_FILENAME, + ): + if candidate.is_file(): + return candidate + return None diff --git a/bergson/config/config.py b/bergson/config/config.py index c86d004d..b11a117e 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -352,10 +352,8 @@ class TrainingConfig(AttributionConfig, Serializable): train_mode: bool = False """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.""" + save_optimizer_state: Literal["none", "last", "all"] = "none" + """Which checkpoints' optimizer second moments to persist. AdamW only.""" weight_decay: float = 0.01 """Weight decay coefficient for AdamW and Muon.""" @@ -381,6 +379,13 @@ 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__() + + # 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" + @dataclass class MetasmoothnessConfig(TrainingConfig): @@ -786,6 +791,11 @@ class ApproxUnrollingConfig(Serializable): """Per-segment optimizer step count; length must == segments. If empty, inferred from per-checkpoint step counts.""" + momentum: float | None = None + """SGD heavy-ball momentum beta; scales lr*steps by 1/(1-beta) (Bae et al. + 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 79970b05..9f0e5979 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,9 @@ def worker( ckpts_path = os.path.join(run_cfg.run_path, "checkpoints") resume = run_cfg.resume + if global_rank == 0: + write_lr_history(ckpts_path, schedule, len(stream)) + fwd_state = trainer.train( fwd_state, stream, @@ -249,8 +253,17 @@ 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", "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/bergson/magic/trainer.py b/bergson/magic/trainer.py index 3ce3681e..6a92c33a 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,9 @@ 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``, + tagged with the step and these hyperparameters. AdamW only. Returns: The final trainer state after training. @@ -615,6 +620,18 @@ def train( pending_save.result() p = os.path.join(save_dir, f"step_{i}.ckpt") + + # 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( + 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/bergson/utils/trainer_export.py b/bergson/utils/trainer_export.py new file mode 100644 index 00000000..93b5d932 --- /dev/null +++ b/bergson/utils/trainer_export.py @@ -0,0 +1,110 @@ +"""Export a run's DCP ``step_.ckpt`` checkpoints to HF-compatible +``checkpoint-/`` model dirs.""" + +import shutil +from pathlib import Path + +import torch +from transformers import AutoTokenizer + +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 + +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__) + + +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 ``step_.ckpt`` to ``checkpoint-/`` HF dirs, in step order. + + ``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 + 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 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) + + 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)) + + 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) + + del trainer, state, model + return exported diff --git a/tests/test_source_trainer_integration.py b/tests/test_source_trainer_integration.py new file mode 100644 index 00000000..f562fef2 --- /dev/null +++ b/tests/test_source_trainer_integration.py @@ -0,0 +1,506 @@ +"""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, + 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" + ) + ckpt = run / "exported" / "checkpoint-0" + ckpt.mkdir(parents=True) + cfg = ApproxUnrollingConfig(checkpoints=[str(ckpt)]) + + 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) + 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) + 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") + ckpt = run / "checkpoint-0" + ckpt.mkdir() + cfg = ApproxUnrollingConfig(checkpoints=[str(ckpt)], model_path="gpt2") + assert resolve(cfg).model_path == "gpt2" + + +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) + + +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( + 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), + ] + + +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) + + +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) + + +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.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.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 + + +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)])) + + +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