feat: seamless Bergson Trainer -> SOURCE attribute

bergson/approx_unrolling/approx_unrolling_math.py

@@ -24,48 +24,66 @@
 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-<N>`` directory
-    (what HF Trainer writes natively) or a bare ``<N>`` step directory.
+    (HF Trainer), a ``step_<N>[.ckpt]`` directory (Bergson trainer), or a
+    bare ``<N>`` 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-<N>' or a bare '<N>' step directory. Set both "
+        "ending in 'checkpoint-<N>', 'step_<N>[.ckpt]', or a bare '<N>' 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."
     )
 
 
 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)
         )

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:

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-<N>`` 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 -- ``<run>/exported/checkpoint-N``
+    and ``<run>/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

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
-    ``<run_path>/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."""
 

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,

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)