Move files around

Author: luciaquirke

README.md

@@ -5,7 +5,7 @@ Data attribution methods estimate the effect on a behavior of interest of removi
 
 ## Core features
 
-Per-token and per-sequence attribution is available everywhere. On-disk gradient stores and on-the-fly queries are supported. Almost every feature is available through both the CLI and a programmatic interface, which use a shared set of configuration dataclasses. Configuration dataclasses are always serialized to disk so commands can be reproduced in one line. To understand every available configuration option, [check out the documentation](https://bergson.readthedocs.io/en/latest/api.html#bergson.IndexConfig). 
+Per-token and per-sequence attribution is available everywhere. On-disk gradient stores and on-the-fly queries are supported. Almost every feature is available through both the CLI and a programmatic interface, which use a shared set of configuration dataclasses. Configuration dataclasses are always serialized to disk so commands can be reproduced in one line. To understand every available configuration option, [check out the documentation](https://bergson.readthedocs.io/en/latest/api.html#bergson.IndexConfig).
 
 Bergson uses FSDP2 or SimpleFSDP, BitsAndBytes, and low-level performance optimizations to support large models, datasets, and clusters. Bergson integrates with HuggingFace Transformers and Datasets, and also supports on-disk datasets in a variety of formats.
 

bergson/__init__.py

@@ -7,7 +7,7 @@
 from .collector.collector import CollectorComputer
 from .collector.gradient_collectors import GradientCollector
 from .collector.in_memory_collector import InMemoryCollector
-from .config import (
+from .config.config import (
     AttentionConfig,
     DataConfig,
     IndexConfig,

bergson/__main__.py

@@ -5,6 +5,8 @@
 
 from simple_parsing import ArgumentParser, ConflictResolution
 
+from bergson.config.config_io import parse_steps, read_config, save_pipeline_config
+
 from .cli.commands import (
     ApproxUnrolling,
     Build,
@@ -20,7 +22,6 @@
     Train,
     Validate,
 )
-from .cli.config_runner import run_config
 
 
 @dataclass
@@ -48,6 +49,31 @@ def execute(self):
         self.command.execute()
 
 
+def run_config(config_path: str, command_registry: dict[str, type]) -> None:
+    """Execute each step of a bergson config YAML in order.
+
+    A fully resolved version of any multi-step config (including default values)
+    is written to ``run_path`` (auto-named under ``runs/`` if not given).
+    Each step also writes its own component ``config.yaml`` into its run directory.
+    """
+    config = read_config(config_path)
+
+    steps = parse_steps(config["steps"], command_registry)
+
+    multi = len(steps) > 1
+
+    if multi:
+        # Optional top-level run path for a multi-step pipeline
+        run_path = config.get("run_path")
+        
+        save_pipeline_config(steps, run_path)
+
+    for i, (cmd_name, cmd) in enumerate(steps, start=1):
+        if multi:
+            print(f"\n[pipeline] step {i}/{len(steps)}: {cmd_name}")
+        cmd.execute()
+
+
 def main():
     """Parse CLI arguments and dispatch to the selected subcommand.
 
@@ -66,9 +92,8 @@ def main():
     command_classes = get_args(Main.__dataclass_fields__["command"].type)
     command_registry = {cls.__name__.lower(): cls for cls in command_classes}
 
-    # Config-file mode: the YAML file is self-describing, so no command verb is
-    # needed. Accept it as the sole argument; leading command words (e.g.
-    # `bergson build run/config.yaml`) are ignored.
+    # Config-file mode: accept a YAML file as the sole argument.
+    # Leading command words (e.g. `bergson build run/config.yaml`) are ignored.
     config_path: str | None = None
     if len(args) == 1 and os.path.isfile(args[0]):
         config_path = args[0]
@@ -79,7 +104,7 @@ def main():
         run_config(config_path, command_registry)
         return
 
-    # CLI-flag mode: standard argparse-style flag parsing.
+    # CLI-flag mode: argparse-style flag parsing.
     parser = ArgumentParser(conflict_resolution=ConflictResolution.EXPLICIT)
     parser.add_arguments(Main, dest="prog")
     prog: Main = parser.parse_args().prog

bergson/approx_unrolling/approx_unrolling_math.py

@@ -9,16 +9,16 @@
 import torch
 from torch import Tensor
 
-from bergson.config import (
+from bergson.cli.commands import Score
+from bergson.config.config import (
     ApproxUnrollingConfig,
     DistributedConfig,
     IndexConfig,
     PreprocessConfig,
     ScoreConfig,
 )
+from bergson.config.config_io import save_run_config
 from bergson.data import load_scores
-from bergson.cli.commands import Score
-from bergson.cli.config_io import save_run_config
 from bergson.distributed import init_dist, launch_distributed_run
 from bergson.hessians.apply_hessian import EkfacApplicator, EkfacConfig
 from bergson.score.score import score_dataset

bergson/approx_unrolling/pipeline.py

@@ -30,13 +30,13 @@
 
 from ..build import build
 from ..cli.commands import Build
-from ..cli.config_io import save_run_config
 from ..config import (
     ApproxUnrollingConfig,
     HessianConfig,
     IndexConfig,
     PreprocessConfig,
 )
+from ..config.config_io import save_run_config
 from ..utils.logger import get_logger
 from .approx_unrolling_math import (
     compute_lr_times_steps_per_segment,

bergson/approx_unrolling/precompute_checkpoints.py

@@ -8,7 +8,7 @@
     CollectorComputer,
     fwd_bwd_hessian_factory,
 )
-from bergson.config import (
+from bergson.config.config import (
     ApproxUnrollingConfig,
     HessianConfig,
     IndexConfig,

bergson/build.py

@@ -8,7 +8,7 @@
 from tqdm.auto import tqdm
 
 from bergson.collection import collect_gradients
-from bergson.config import HessianConfig, IndexConfig, PreprocessConfig
+from bergson.config.config import HessianConfig, IndexConfig, PreprocessConfig
 from bergson.data import allocate_batches
 from bergson.distributed import (
     cap_world_size_to_dataset,

bergson/builder.py

@@ -5,7 +5,7 @@
 import torch.distributed as dist
 from datasets import Dataset
 
-from .config import PreprocessConfig
+from .config.config import PreprocessConfig
 from .data import compute_num_token_grads, create_index, create_token_index
 from .process_grads import (
     get_trackstar_hessian,

bergson/cli/commands.py

@@ -1,17 +1,17 @@
 """CLI command definitions.
 
 Each command is a thin dataclass that validates, persists its ``config.yaml``,
-and dispatches. They live here rather than in ``__main__`` so pipelines can 
-lazily import lower-level commands for step serialization without importing 
-the corresponding CLI entrypoint.
+and dispatches. They live here rather than in ``__main__`` so pipelines can
+import lower-level commands for config serialization without importing the 
+corresponding CLI entrypoint.
 """
 
 from dataclasses import dataclass
 
 from simple_parsing import Serializable
 
 from ..build import build
-from ..config import (
+from ..config.config import (
     ApproxUnrollingConfig,
     HessianConfig,
     HessianPipelineConfig,
@@ -24,14 +24,14 @@
     TrainingConfig,
     ValidationConfig,
 )
+from ..config.config_io import save_run_config
 from ..diagnose import DiagnoseConfig, diagnose
 from ..hessians.hessian_approximations import approximate_hessians
 from ..magic import MagicConfig, run_magic
 from ..process_grads import mix_autocorrelation_matrices
 from ..query.query_index import query
 from ..score.score import score_dataset
 from ..utils.worker_utils import validate_run_path
-from .config_io import save_run_config
 
 
 @dataclass
@@ -82,7 +82,10 @@ def execute(self):
                 "if skip_index is True HessianConfig.method must be provided"
             )
 
-        if self.hessian_cfg is not None and self.hessian_cfg.method != "autocorrelation":
+        if (
+            self.hessian_cfg is not None
+            and self.hessian_cfg.method != "autocorrelation"
+        ):
             raise ValueError(
                 f"build only supports autocorrelation Hessians, got "
                 f"'{self.hessian_cfg.method}'. Use the `hessian` command for "
@@ -224,7 +227,6 @@ def execute(self):
         score_dataset(self.index_cfg, self.score_cfg, self.preprocess_cfg)
 
 
-
 @dataclass
 class Trackstar(Serializable):
     """Run hessians, build, and score as a single pipeline."""
@@ -234,7 +236,7 @@ class Trackstar(Serializable):
     trackstar_cfg: TrackstarConfig
 
     def execute(self):
-        from ..trackstar import trackstar
+        from .trackstar import trackstar
 
         save_run_config(self, self.index_cfg.run_path)
         trackstar(self.index_cfg, self.trackstar_cfg)