diff --git a/bergson/config/config.py b/bergson/config/config.py index 69cff426..40f64ee3 100644 --- a/bergson/config/config.py +++ b/bergson/config/config.py @@ -454,9 +454,6 @@ class ValidationConfig(TrainingConfig, ABC): num_subsets: int = 100 """Number of leave-k-out subsets for Spearman correlation.""" - subset_strategy: Literal["random"] = "random" - """Strategy for selecting leave-k-out subsets for validation.""" - subset_weight: float = 0.0 """Training weight assigned to each subset's documents during the retrain (the rest stay at 1.0). ``0.0`` (default) is standard leave-k-out removal.""" diff --git a/bergson/magic/config.py b/bergson/magic/config.py index 2b238609..33bc4d8b 100644 --- a/bergson/magic/config.py +++ b/bergson/magic/config.py @@ -17,9 +17,9 @@ class MagicConfig(ValidationConfig): """Whether to compute attribution scores per token (instead of per sequence); the same toggle as ``IndexConfig.attribute_tokens``.""" - skip_validation: bool = False - """Stop after computing and saving attribution scores, before the - leave-k-out retraining loop. Useful for score-only MAGIC runs.""" + skip_validation: bool = True + """Set to False to run a leave-k-out retraining validation loop in the + same job.""" # TODO(Lucia Quirke, December 2026): remove per_token backward compatibility. per_token: bool = False diff --git a/bergson/validate.py b/bergson/validate.py index efb8e290..b84f8288 100644 --- a/bergson/validate.py +++ b/bergson/validate.py @@ -324,7 +324,7 @@ def validate_scores( if os.path.exists(subsets_path): with open(subsets_path) as f: subsets = [torch.tensor(s, dtype=torch.long) for s in json.load(f)] - elif run_cfg.subset_strategy == "random": + else: rng = torch.Generator().manual_seed(run_cfg.seed) if run_cfg.subset_fraction > 0: # Draw potentially overlapping samples @@ -348,8 +348,6 @@ def validate_scores( subsets = list(perm.chunk(run_cfg.num_subsets)) rng = random.Random(run_cfg.seed) rng.shuffle(subsets) - else: - raise ValueError(f"Unknown subset strategy: {run_cfg.subset_strategy}") csv_path = os.path.join(run_cfg.run_path, "validation.csv") val_csv_writer = CSVWriter( diff --git a/docs/pipeline.rst b/docs/building_blocks.rst similarity index 99% rename from docs/pipeline.rst rename to docs/building_blocks.rst index e6b74df7..faa14ee5 100644 --- a/docs/pipeline.rst +++ b/docs/building_blocks.rst @@ -1,5 +1,5 @@ -Pipeline Concepts -================= +Building Blocks +=============== Bergson's post-hoc attribution exposes three generic building blocks — ``build``, ``reduce``, and ``score`` — that together implement gradient-based data attribution. This page diff --git a/docs/cli.rst b/docs/cli.rst index 5ab8db32..942fcaa3 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -149,7 +149,7 @@ Method Pipelines --method kfac \ --hessian_cfg.ev_correction true -See ``examples/magic/compare/q3_ekfac.yaml`` for a complete pipeline +See ``examples/magic/compare_wikitext/ekfac.yaml`` for a complete pipeline configuration. .. autoclass:: bergson.__main__.ApproxUnrolling diff --git a/docs/magic.rst b/docs/magic.rst index 1de8aa6d..fd845489 100644 --- a/docs/magic.rst +++ b/docs/magic.rst @@ -32,6 +32,12 @@ Output files After a run completes, ``run_cfg.run_path`` contains: +* ``config.yaml`` — the serialized configuration needed to replicate the run. +* ``checkpoints/`` — forward-pass trainer checkpoints, plus + ``log_history.json`` (the per-step learning rates the schedule produced). + With the default ``cleanup_ckpts=True`` the backward pass deletes + checkpoints as it consumes them, but multi-query runs keep them for re-use. + * ``scores/`` — a score directory, the same self-describing format the scoring pipeline writes. ``info.json`` records ``attribute_tokens`` and ``num_scores``, so consumers never infer the layout from the shape. @@ -50,10 +56,15 @@ After a run completes, ``run_cfg.run_path`` contains: ``length - 1`` values, the positions ``weighted_causal_lm_ce`` can reach — and are unpacked back into the dense grid on load. +* ``per_query/q{i}.pt`` — per-query runs only. The score tensor for query + document ``i``, written as soon as that query's backward finishes so an + interrupted run resumes without redoing completed queries. The trailing + query axis in ``scores/`` is these tensors stacked. + * ``scores/doc_ids.npy`` — written for every per-token run, shape - ``(num_chunks, seq_len)`` matching the loaded scores row-for-row. Each - entry is the original (pre-shuffle) document id for that token position. - Downstream aggregation is one line: + ``(num_chunks, seq_len)`` matching the loaded scores. Each + entry is the original (pre-shuffle) document id for that token position, + so the scores can be aggregated over chunks: .. code-block:: python @@ -62,8 +73,15 @@ After a run completes, ``run_cfg.run_path`` contains: scores, _ = load_scores_loss_signed("runs/magic/scores") doc_ids = torch.from_numpy(np.load("runs/magic/scores/doc_ids.npy")) num_docs = int(doc_ids.max()) + 1 - per_doc = torch.zeros(num_docs, dtype=scores.dtype) - per_doc.scatter_add_(0, doc_ids.flatten(), scores.flatten()) + + # Trailing axis is the query axis on per-query runs, absent otherwise; + # reshaping to it keeps both cases on one path. + flat = scores.reshape(doc_ids.numel(), -1) + per_doc = torch.zeros(num_docs, flat.shape[1], dtype=flat.dtype) + per_doc.scatter_add_(0, doc_ids.flatten()[:, None].expand_as(flat), flat) + + ``per_doc`` comes back as ``(num_docs, num_query_docs)``, or + ``(num_docs, 1)`` for a single-query run. When ``data.chunk_length > 0`` the ``doc_ids`` column comes from ``tokenize_and_chunk`` and chunks may pack multiple docs or split one @@ -72,14 +90,23 @@ After a run completes, ``run_cfg.run_path`` contains: past the row's actual length carry zero MAGIC score and contribute nothing to the scatter-add. -* ``config.yaml`` — serialized ``MagicConfig`` used for the run. -* ``validation.csv`` — leave-subset-out validation results (if validation - was run). +Models and optimizer state +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +* ``optimizer.pt`` — second moments of the trained optimizer state, when + ``save_optimizer_state`` is not ``"none"``. +* ``retrained/base/`` and ``retrained/subset_{i}/`` — the fully-trained + model and each leave-subset-out retrain, with tokenizer, when + ``save_models=True``. ``retrained/base`` is the query baseline + ``evaluate_retrained`` reads; the ``subset_{i}`` retrains only exist for + a run that actually validates (``bergson validate``, or ``bergson magic + --skip_validation False``). (Commands outside the leave-k-out family — + plain ``bergson train`` — write the trained model to ``model/`` instead.) -Metasmoothness ---------------- +Smoothness +---------- -MAGIC is valid when the function you are differentiating through is metasmooth. There a few heuristics known to encourage metasmoothness: +MAGIC is valid when the model training function you differentiate through is smooth with respect to the data weightings (metasmooth). There a few heuristics known to encourage smoothness: * Use the Muon optimizer * Increase batch size @@ -89,7 +116,7 @@ MAGIC is valid when the function you are differentiating through is metasmooth. * QK norm * Tune weight decay -Many of these methods boil down to "Identify and manage spikes in your training loss." You can measure your metasmoothness with ``bergson metasmoothness``. +Many of these methods boil down to "Identify and manage spikes in your training loss." You can measure your smoothness with ``bergson metasmoothness``. Core components ^^^^^^^^^^^^^^^ diff --git a/examples/compare_wikitext/README.md b/examples/compare_wikitext/README.md new file mode 100644 index 00000000..ed587d71 --- /dev/null +++ b/examples/compare_wikitext/README.md @@ -0,0 +1,8 @@ +| method | mean ρ | median ρ | min | max | queries p<.05 | +|---|---|---|---|---|---| +| MAGIC (per-query) | 0.957 | 0.963 | 0.880 | 0.984 | 50/50 | +| EK-FAC | 0.470 | 0.461 | 0.145 | 0.751 | 46/50 | +| SOURCE-Adam | 0.243 | 0.253 | -0.038 | 0.434 | 30/50 | +| TrackStar | 0.241 | 0.239 | -0.037 | 0.454 | 36/50 | +| SOURCE | 0.221 | 0.212 | -0.089 | 0.421 | 27/50 | +| TrackStar+Adam | 0.189 | 0.190 | -0.107 | 0.471 | 23/50 | diff --git a/examples/compare_wikitext/ekfac.yaml b/examples/compare_wikitext/ekfac.yaml new file mode 100644 index 00000000..a2736f2e --- /dev/null +++ b/examples/compare_wikitext/ekfac.yaml @@ -0,0 +1,73 @@ +# Pipeline: compute EK-FAC scores for the magic-trained model against +# test[0:50] (one score column per query), then run leave-k-out validation on +# those scores against the retrained models from magic.yaml. + +# Run with `bergson examples/compare_wikitext/ekfac.yaml` + +run_path: runs/compare_wikitext/ekfac +steps: + - ekfac: + index_cfg: + run_path: runs/compare_wikitext/ekfac + model: runs/compare_wikitext/random/retrained/base + overwrite: true + # GPT-2's max context length is 1024. + token_batch_size: 1024 + # Cap batch size in docs so the bin-packer doesn't produce + # large batches of very few tokens. This matters for EK-FAC + # because gradients aren't compressed. + max_batch_size: 64 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + hessian_cfg: + method: kfac + ev_correction: true + + score_cfg: + batch_size: 1024 + query_batch_size: 64 + + preprocess_cfg: + unit_normalize: false + + hessian_pipeline_cfg: + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + # One score column per query, for per-query LDS. + query_aggregation: none + inversion_cfg: + damping_factor: 0.1 + + - validate: + run_path: runs/compare_wikitext/validate_ekfac_random + model: gpt2 + overwrite: true + + scores: runs/compare_wikitext/ekfac/scores + retrained_dir: runs/compare_wikitext/random + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + batch_size: 64 diff --git a/examples/compare_wikitext/magic.yaml b/examples/compare_wikitext/magic.yaml new file mode 100644 index 00000000..874f32b7 --- /dev/null +++ b/examples/compare_wikitext/magic.yaml @@ -0,0 +1,54 @@ +# Train GPT-2 on WikiText, computing per-query MAGIC scores (one backward per +# query, 50 queries) and leave-k-out validation over 100 random 1%-drop +# subsets. bs 256 / 4 epochs / eps_root 1e-8 scored per-query LDS 0.93-0.98 +# in the 2026-08-03 batch-size sweep (runs/bs_eps1e8). The retrained models +# are saved so every other yaml in this directory validates against the same +# family, and the final Adam second moments are kept for trackstar_adam.yaml. + +# Run with `bergson examples/compare_wikitext/magic.yaml` + +run_path: runs/compare_wikitext/magic +steps: + - magic: + run_path: runs/compare_wikitext/random + model: gpt2 + overwrite: true + # Resume: keep finished per_query/q.pt scores and fast-forward + # training from the last checkpoint after an interruption. + resume: true + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + batch_size: 256 + # Bound double-backward memory; no trajectory change without dropout. + # Without the cap the retained graph OOMs 48GB GPUs (vocab-sized + # tensors per sequence). + grad_accum_steps: 2 + double_backward_batch_size: 8 + num_epochs: 4 + lr_schedule: + lr_scheduler_type: polynomial + lr: 0.0008 + lr_start: 1e-6 + lr_end: 0.00008 + warmup_steps: 0.25 + + skip_validation: false + num_subsets: 100 + subset_fraction: 0.01 + + save_models: true + save_optimizer_state: last + wandb_project: magic diff --git a/examples/compare_wikitext/source.yaml b/examples/compare_wikitext/source.yaml new file mode 100644 index 00000000..66ff1109 --- /dev/null +++ b/examples/compare_wikitext/source.yaml @@ -0,0 +1,102 @@ +# SOURCE (approximate unrolling, Bae et al. 2024) scores for the +# magic-trained model against test[0:50], validated against the retrained +# models from magic.yaml. Plain variant: no optimizer preconditioning; see +# source_adam.yaml for the Adam variant. Shares its evenly spaced interval +# checkpoints (steps 12..72; sqrt mode never saves the final step) — the +# train step resumes instantly when runs/compare_wikitext/interval already +# exists. + +# Run with `bergson examples/compare_wikitext/source.yaml` + +run_path: runs/compare_wikitext/source +steps: + - train: + run_path: runs/compare_wikitext/interval + model: gpt2 + resume: true + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + batch_size: 256 + grad_accum_steps: 2 + num_epochs: 4 + lr_schedule: + lr_scheduler_type: polynomial + lr: 0.0008 + lr_start: 1e-6 + lr_end: 0.00008 + warmup_steps: 0.25 + + save_mode: interval + save_interval: 12 + save_optimizer_state: all + + - approxunrolling: + index_cfg: + run_path: runs/compare_wikitext/source + model: runs/compare_wikitext/random/retrained/base + overwrite: true + precision: fp32 + # GPT-2's max context length is 1024. + token_batch_size: 1024 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + approx_unrolling_cfg: + checkpoints: + - runs/compare_wikitext/interval/checkpoints/step_12.ckpt + - runs/compare_wikitext/interval/checkpoints/step_24.ckpt + - runs/compare_wikitext/interval/checkpoints/step_36.ckpt + - runs/compare_wikitext/interval/checkpoints/step_48.ckpt + - runs/compare_wikitext/interval/checkpoints/step_60.ckpt + - runs/compare_wikitext/interval/checkpoints/step_72.ckpt + segments: 3 + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + # One score column per query, for per-query LDS. + query_aggregation: none + + - validate: + run_path: runs/compare_wikitext/validate_source_random + model: gpt2 + overwrite: true + + scores: runs/compare_wikitext/source/scores + retrained_dir: runs/compare_wikitext/random + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + batch_size: 64 diff --git a/examples/compare_wikitext/source_adam.yaml b/examples/compare_wikitext/source_adam.yaml new file mode 100644 index 00000000..36eaa42a --- /dev/null +++ b/examples/compare_wikitext/source_adam.yaml @@ -0,0 +1,109 @@ +# Adam-preconditioned SOURCE (approximate unrolling, Bae et al. 2024, +# App. C) scores for the magic-trained model against test[0:50], validated +# against the retrained models from magic.yaml. +# +# MAGIC needs sqrt-mode checkpoints, which are unevenly spaced and never +# include the final step, so the first step re-runs the identical training +# recipe (same seed and batch schedule) with save_mode interval to get +# evenly spaced checkpoints including the final state, each with its Adam +# second moments. 4 epochs x 18 steps = 72 steps; interval 12 saves +# step_{0,12,24,36,48,60,72}. Raw DCP checkpoints are auto-exported to +# runs/compare_wikitext/interval/exported/; lr_list and step_size_list are +# inferred from the run's log_history.json. + +# Run with `bergson examples/compare_wikitext/source_adam.yaml` + +run_path: runs/compare_wikitext/source_adam +steps: + - train: + run_path: runs/compare_wikitext/interval + model: gpt2 + # Shared with source.yaml; resumes instantly when the run exists. + resume: true + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + batch_size: 256 + grad_accum_steps: 2 + num_epochs: 4 + lr_schedule: + lr_scheduler_type: polynomial + lr: 0.0008 + lr_start: 1e-6 + lr_end: 0.00008 + warmup_steps: 0.25 + + save_mode: interval + save_interval: 12 + save_optimizer_state: all + + - approxunrolling: + index_cfg: + run_path: runs/compare_wikitext/source_adam + model: runs/compare_wikitext/random/retrained/base + overwrite: true + precision: fp32 + # GPT-2's max context length is 1024. + token_batch_size: 1024 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + hessian_cfg: + method: kfac + hessian_dtype: fp32 + ev_correction: true + + approx_unrolling_cfg: + checkpoints: + - runs/compare_wikitext/interval/checkpoints/step_12.ckpt + - runs/compare_wikitext/interval/checkpoints/step_24.ckpt + - runs/compare_wikitext/interval/checkpoints/step_36.ckpt + - runs/compare_wikitext/interval/checkpoints/step_48.ckpt + - runs/compare_wikitext/interval/checkpoints/step_60.ckpt + - runs/compare_wikitext/interval/checkpoints/step_72.ckpt + segments: 3 + use_adam_preconditioner: true + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + # One score column per query, for per-query LDS. + query_aggregation: none + + - validate: + run_path: runs/compare_wikitext/validate_source_adam_random + model: gpt2 + overwrite: true + + scores: runs/compare_wikitext/source_adam/scores + retrained_dir: runs/compare_wikitext/random + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + batch_size: 64 diff --git a/examples/compare_wikitext/trackstar.yaml b/examples/compare_wikitext/trackstar.yaml new file mode 100644 index 00000000..1e8ecc23 --- /dev/null +++ b/examples/compare_wikitext/trackstar.yaml @@ -0,0 +1,63 @@ +# Compute TrackStar attribution scores for the magic-trained model against +# test[0:50] (one score column per query), then run leave-k-out validation on +# those scores against the retrained models from magic.yaml. No optimizer +# normalization; see trackstar_adam.yaml for the Adam second-moment-normalized +# variant. + +# Run with `bergson examples/compare_wikitext/trackstar.yaml` + +run_path: runs/compare_wikitext/trackstar +steps: + - trackstar: + index_cfg: + projection_dim: 32 + run_path: runs/compare_wikitext/trackstar + model: runs/compare_wikitext/random/retrained/base + overwrite: true + # GPT-2's max context length is 1024 + token_batch_size: 1024 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + trackstar_cfg: + stats_sample_size: 10000 + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + preprocess_cfg: + unit_normalize: true + aggregation: none + score_cfg: + batch_size: 1024 + + - validate: + run_path: runs/compare_wikitext/validate_trackstar_random + model: gpt2 + overwrite: true + + scores: runs/compare_wikitext/trackstar/scores + retrained_dir: runs/compare_wikitext/random + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + batch_size: 64 diff --git a/examples/compare_wikitext/trackstar_adam.yaml b/examples/compare_wikitext/trackstar_adam.yaml new file mode 100644 index 00000000..3608acbb --- /dev/null +++ b/examples/compare_wikitext/trackstar_adam.yaml @@ -0,0 +1,62 @@ +# TrackStar with Adam second-moment gradient normalization (Chang et al. +# 2024), using the optimizer state saved at the end of the magic.yaml +# training run. Otherwise identical to trackstar.yaml. + +# Run with `bergson examples/compare_wikitext/trackstar_adam.yaml` + +run_path: runs/compare_wikitext/trackstar_adam +steps: + - trackstar: + index_cfg: + projection_dim: 32 + run_path: runs/compare_wikitext/trackstar_adam + model: runs/compare_wikitext/random/retrained/base + optimizer_state: runs/compare_wikitext/random/optimizer.pt + overwrite: true + # GPT-2's max context length is 1024 + token_batch_size: 1024 + # GPT-2 uses an nn.Linear for its LM head. + filter_modules: "lm_head" + distributed: + nproc_per_node: 8 + nnode: 1 + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + trackstar_cfg: + stats_sample_size: 10000 + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + preprocess_cfg: + unit_normalize: true + aggregation: none + score_cfg: + batch_size: 1024 + + - validate: + run_path: runs/compare_wikitext/validate_trackstar_adam_random + model: gpt2 + overwrite: true + + scores: runs/compare_wikitext/trackstar_adam/scores + retrained_dir: runs/compare_wikitext/random + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:50]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + batch_size: 64 diff --git a/examples/magic/compare/q3_ekfac.yaml b/examples/magic/compare/q3_ekfac.yaml deleted file mode 100644 index 7ad59932..00000000 --- a/examples/magic/compare/q3_ekfac.yaml +++ /dev/null @@ -1,189 +0,0 @@ -# Pipeline: compute EK-FAC scores for the magic-trained model from q3_random -# against test[3:4], then run leave-k-out validation on those scores. - -# Run with `bergson examples/magic/compare/q3_ekfac.yaml` - -run_path: runs/compare/q3_ekfac -steps: - - ekfac: - index_cfg: - run_path: runs/compare/q3_ekfac - model: runs/compare/q3_random/hf_model - overwrite: true - # GPT-2's max context length is 1024. - token_batch_size: 1024 - # Cap batch size in docs so the bin-packer doesn't produce - # large batches of very few tokens. This matters for EK-FAC - # because gradients aren't compressed. - max_batch_size: 64 - # GPT-2 uses an nn.Linear for its LM head. - filter_modules: "lm_head" - distributed: - nproc_per_node: 4 - nnode: 4 - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - truncation: true - - hessian_cfg: - method: kfac - ev_correction: true - - score_cfg: - batch_size: 1024 - - preprocess_cfg: - unit_normalize: false - - hessian_pipeline_cfg: - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - inversion_cfg: - damping_factor: 0.1 - - - validate: - run_path: runs/compare/q3_validate_ekfac_random - model: gpt2 - overwrite: true - - scores: runs/compare/q3_ekfac/scores - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: random - wandb_project: magic - - - validate: - run_path: runs/compare/q3_validate_ekfac_sorted - model: gpt2 - overwrite: true - - scores: runs/compare/q3_ekfac/scores - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: sorted - wandb_project: magic - - - validate: - run_path: runs/compare/q3_validate_ekfac_random_exclude_0 - model: gpt2 - overwrite: true - - scores: runs/compare/q3_ekfac/scores - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: random - exclude_zero_scores: true - wandb_project: magic - - - validate: - run_path: runs/compare/q3_validate_ekfac_sorted_exclude_0 - model: gpt2 - overwrite: true - - scores: runs/compare/q3_ekfac/scores - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: sorted - exclude_zero_scores: true - wandb_project: magic diff --git a/examples/magic/compare/q3_magic.yaml b/examples/magic/compare/q3_magic.yaml deleted file mode 100644 index 9d068cbd..00000000 --- a/examples/magic/compare/q3_magic.yaml +++ /dev/null @@ -1,146 +0,0 @@ -# Compute MAGIC attribution scores with random subset validation -# then run a sorted-subset validate against the same scores. - -# Run with `bergson examples/magic/compare/q3_magic.yaml` - -run_path: runs/compare/q3_magic -steps: - - magic: - run_path: runs/compare/q3_random - model: gpt2 - overwrite: true - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: random - wandb_project: magic - - - validate: - run_path: runs/compare/q3_sorted - model: gpt2 - overwrite: true - - scores: runs/compare/q3_random/scores.pt - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: sorted - wandb_project: magic - - - validate: - run_path: runs/compare/q3_validate_magic_random_exclude_0 - model: gpt2 - overwrite: true - - scores: runs/compare/q3_random/scores.pt - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: random - exclude_zero_scores: true - wandb_project: magic - - - validate: - run_path: runs/compare/q3_validate_magic_sorted_exclude_0 - model: gpt2 - overwrite: true - - scores: runs/compare/q3_random/scores.pt - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: sorted - exclude_zero_scores: true - wandb_project: magic diff --git a/examples/magic/compare/q3_trackstar.yaml b/examples/magic/compare/q3_trackstar.yaml deleted file mode 100644 index 4dce4a0c..00000000 --- a/examples/magic/compare/q3_trackstar.yaml +++ /dev/null @@ -1,183 +0,0 @@ -# Compute TrackStar attribution scores for a model generated -# using q3_magic.yaml using test[3:4] as the query, then run leave-k-out -# validation on those scores. - -# Run with `bergson examples/magic/compare/q3_trackstar.yaml` - -run_path: runs/compare/q3_trackstar -steps: - - trackstar: - index_cfg: - projection_dim: 16 - run_path: runs/compare/q3_trackstar - model: runs/compare/q3_random/hf_model - overwrite: true - # GPT-2's max context length is 1024 - token_batch_size: 1024 - # GPT-2 uses an nn.Linear for its LM head. - filter_modules: "lm_head" - distributed: - nproc_per_node: 4 - nnode: 4 - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 0 - truncation: true - - trackstar_cfg: - stats_sample_size: 10000 - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - truncation: true - preprocess_cfg: - unit_normalize: true - aggregation: mean - score_cfg: - batch_size: 1024 - - - validate: - run_path: runs/compare/q3_validate_trackstar_random - model: gpt2 - overwrite: true - - scores: runs/compare/q3_trackstar/scores - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: random - wandb_project: magic - - - validate: - run_path: runs/compare/q3_validate_trackstar_sorted - model: gpt2 - overwrite: true - - scores: runs/compare/q3_trackstar/scores - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: sorted - wandb_project: magic - - - validate: - run_path: runs/compare/q3_validate_trackstar_random_exclude_0 - model: gpt2 - overwrite: true - - scores: runs/compare/q3_trackstar/scores - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: random - exclude_zero_scores: true - wandb_project: magic - - - validate: - run_path: runs/compare/q3_validate_trackstar_sorted_exclude_0 - model: gpt2 - overwrite: true - - scores: runs/compare/q3_trackstar/scores - - data: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "train" - chunk_length: 512 - - query: - dataset: Salesforce/wikitext - subset: wikitext-2-raw-v1 - split: "test[3:4]" - chunk_length: 0 - - distributed: - nproc_per_node: 4 - nnode: 4 - - batch_size: 256 - num_epochs: 2 - lr_schedule: - lr_scheduler_type: polynomial - lr: 0.0008 - lr_start: 1e-6 - lr_end: 0.00008 - warmup_steps: 0.25 - - subset_strategy: sorted - exclude_zero_scores: true - wandb_project: magic diff --git a/examples/magic/gpt2_wikitext.yaml b/examples/magic/gpt2_wikitext_replication.yaml similarity index 77% rename from examples/magic/gpt2_wikitext.yaml rename to examples/magic/gpt2_wikitext_replication.yaml index c6d7b102..2cf11542 100644 --- a/examples/magic/gpt2_wikitext.yaml +++ b/examples/magic/gpt2_wikitext_replication.yaml @@ -3,6 +3,7 @@ # Hyperparameters are inferred from the MAGIC paper, the metagradients # paper it cites, and the datacomp competition that the metagradients paper cites. +# Original batch size and loss reduction were lost and re-selected empirically. # Per-query MAGIC LDS = 0.952 (95% CI [0.944, 0.959]), measured over m=50 # queries against an N=100 leave-1%-out retrain bank. @@ -11,7 +12,7 @@ steps: - magic: - run_path: runs/gpt2_wikitext + run_path: runs/gpt2_wikitext_replication model: gpt2 overwrite: true cleanup_ckpts: false @@ -20,10 +21,6 @@ steps: adam_beta1: 0.95 adam_beta2: 0.975 eps_root: 1.0e-8 - # The metagradients paper documents sum_of_means for Gemma-IFT, but the - # MAGIC paper's Fig. 5 GPT-2 panel shows 1%-drop retrains moving true - # loss by only ~0.006 — matching our mean-reduction banks (0.007-0.009) - # and excluding sum_of_means (0.2, nearly batch-size-invariant). loss_reduction: mean data: @@ -50,9 +47,10 @@ steps: lr_end: 0.00008 warmup_steps: 0.25 + skip_validation: false num_subsets: 100 subset_fraction: 0.01 wandb_project: magic - save_optimizer_state: True + save_optimizer_state: last save_models: True diff --git a/examples/replicate_bae_approx_unrolling_source/wikitext_gpt2_retrain.yaml b/examples/replicate_bae_approx_unrolling_source/wikitext_gpt2_retrain.yaml index 40b70d5d..d4ffadc9 100644 --- a/examples/replicate_bae_approx_unrolling_source/wikitext_gpt2_retrain.yaml +++ b/examples/replicate_bae_approx_unrolling_source/wikitext_gpt2_retrain.yaml @@ -35,7 +35,6 @@ steps: lr_scheduler_type: constant lr: 3.0e-05 train_mode: true - subset_strategy: random num_subsets: 100 subset_fraction: 0.5 save_models: true