From e4ca451c1e7308f15c7a8b8db09622d642c7fd7d Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Tue, 4 Aug 2026 14:18:42 +0900 Subject: [PATCH 1/7] Clean up MAGIC replication YAML --- .../{gpt2_wikitext.yaml => gpt2_wikitext_replication.yaml} | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) rename examples/magic/{gpt2_wikitext.yaml => gpt2_wikitext_replication.yaml} (79%) diff --git a/examples/magic/gpt2_wikitext.yaml b/examples/magic/gpt2_wikitext_replication.yaml similarity index 79% rename from examples/magic/gpt2_wikitext.yaml rename to examples/magic/gpt2_wikitext_replication.yaml index c6d7b102..3981e69b 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: From e6d3efcf8008ef3516c5da061a4e6d7035fedc4a Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Tue, 4 Aug 2026 09:10:14 +0000 Subject: [PATCH 2/7] clean up yamls/subset flag --- bergson/config/config.py | 3 - bergson/validate.py | 4 +- docs/cli.rst | 2 +- examples/compare_wikitext/ekfac.yaml | 87 ++++++++ examples/compare_wikitext/magic.yaml | 72 +++++++ examples/compare_wikitext/trackstar.yaml | 76 +++++++ examples/magic/compare/q3_ekfac.yaml | 189 ------------------ examples/magic/compare/q3_magic.yaml | 146 -------------- examples/magic/compare/q3_trackstar.yaml | 183 ----------------- examples/magic/gpt2_wikitext_replication.yaml | 2 +- .../wikitext_gpt2_retrain.yaml | 1 - 11 files changed, 238 insertions(+), 527 deletions(-) create mode 100644 examples/compare_wikitext/ekfac.yaml create mode 100644 examples/compare_wikitext/magic.yaml create mode 100644 examples/compare_wikitext/trackstar.yaml delete mode 100644 examples/magic/compare/q3_ekfac.yaml delete mode 100644 examples/magic/compare/q3_magic.yaml delete mode 100644 examples/magic/compare/q3_trackstar.yaml 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/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/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/examples/compare_wikitext/ekfac.yaml b/examples/compare_wikitext/ekfac.yaml new file mode 100644 index 00000000..5a079128 --- /dev/null +++ b/examples/compare_wikitext/ekfac.yaml @@ -0,0 +1,87 @@ +# Pipeline: compute EK-FAC scores for the magic-trained model from q3_random +# against test[0:50] (one score column per query), then run leave-k-out +# validation on those scores, with and without zero-score exclusion. + +# Run with `bergson examples/magic/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: &validate + run_path: runs/compare_wikitext/validate_ekfac_random + model: gpt2 + overwrite: true + + scores: runs/compare_wikitext/ekfac/scores + + 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: 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 + + wandb_project: magic + + - validate: + # YAML merge key: copy the contents of &validate + <<: *validate + run_path: runs/compare_wikitext/validate_ekfac_random_exclude_0 + exclude_zero_scores: true diff --git a/examples/compare_wikitext/magic.yaml b/examples/compare_wikitext/magic.yaml new file mode 100644 index 00000000..a8300c7b --- /dev/null +++ b/examples/compare_wikitext/magic.yaml @@ -0,0 +1,72 @@ +# Compute per-query MAGIC attribution scores (one backward per query, 50 +# queries) with leave-k-out validation, then revalidate the same scores +# excluding zero scores. + +# Run with `bergson examples/magic/compare_wikitext/magic.yaml` + +run_path: runs/compare_wikitext/magic +steps: + - magic: + run_path: runs/compare_wikitext/random + model: gpt2 + overwrite: 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: 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 + + # Export the trained model for q3_ekfac.yaml / q3_trackstar.yaml. + save_models: true + wandb_project: magic + + - validate: + run_path: runs/compare_wikitext/validate_magic_random_exclude_0 + model: gpt2 + overwrite: true + + scores: runs/compare_wikitext/random/scores.pt + + 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 + num_epochs: 2 + lr_schedule: + lr_scheduler_type: polynomial + lr: 0.0008 + lr_start: 1e-6 + lr_end: 0.00008 + warmup_steps: 0.25 + + exclude_zero_scores: true + wandb_project: magic diff --git a/examples/compare_wikitext/trackstar.yaml b/examples/compare_wikitext/trackstar.yaml new file mode 100644 index 00000000..6373d654 --- /dev/null +++ b/examples/compare_wikitext/trackstar.yaml @@ -0,0 +1,76 @@ +# Compute TrackStar attribution scores for a model generated +# using q3_magic.yaml against test[0:50] (one score column per query), then +# run leave-k-out validation on those scores, with and without zero-score +# exclusion. + +# Run with `bergson examples/magic/compare_wikitext/trackstar.yaml` + +run_path: runs/compare_wikitext/trackstar +steps: + - trackstar: + index_cfg: + projection_dim: 16 + 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: &validate + run_path: runs/compare_wikitext/validate_trackstar_random + model: gpt2 + overwrite: true + + scores: runs/compare_wikitext/trackstar/scores + + 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: 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 + + wandb_project: magic + + - validate: + # YAML merge key: copy the contents of &validate + <<: *validate + run_path: runs/compare_wikitext/validate_trackstar_random_exclude_0 + exclude_zero_scores: true 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_replication.yaml b/examples/magic/gpt2_wikitext_replication.yaml index 3981e69b..bf39eccf 100644 --- a/examples/magic/gpt2_wikitext_replication.yaml +++ b/examples/magic/gpt2_wikitext_replication.yaml @@ -51,5 +51,5 @@ steps: 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 From cf9b2ae2458e8b71ed33efe8fad5a1c080413240 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 04:25:42 +0000 Subject: [PATCH 3/7] compare_wikitext: shared retrain family, TrackStar+Adam and SOURCE-Adam, LDS results One magic run builds the 100x1% leave-k-out family (validated LDS 0.957); ekfac/trackstar/trackstar_adam/source_adam score the same model and validate against it via retrained_dir. source_adam re-trains with save_mode interval for evenly spaced checkpoints including the final step. Validate steps use batch_size 64 because evaluate_retrained computes query losses on one GPU. --- examples/compare_wikitext/README.md | 7 ++ examples/compare_wikitext/ekfac.yaml | 32 ++---- examples/compare_wikitext/magic.yaml | 61 ++++------ examples/compare_wikitext/source_adam.yaml | 108 ++++++++++++++++++ examples/compare_wikitext/trackstar.yaml | 37 ++---- examples/compare_wikitext/trackstar_adam.yaml | 62 ++++++++++ 6 files changed, 219 insertions(+), 88 deletions(-) create mode 100644 examples/compare_wikitext/README.md create mode 100644 examples/compare_wikitext/source_adam.yaml create mode 100644 examples/compare_wikitext/trackstar_adam.yaml diff --git a/examples/compare_wikitext/README.md b/examples/compare_wikitext/README.md new file mode 100644 index 00000000..c55d36b9 --- /dev/null +++ b/examples/compare_wikitext/README.md @@ -0,0 +1,7 @@ +| 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 | +| 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 index 5a079128..a2736f2e 100644 --- a/examples/compare_wikitext/ekfac.yaml +++ b/examples/compare_wikitext/ekfac.yaml @@ -1,8 +1,8 @@ -# Pipeline: compute EK-FAC scores for the magic-trained model from q3_random -# against test[0:50] (one score column per query), then run leave-k-out -# validation on those scores, with and without zero-score exclusion. +# 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/magic/compare_wikitext/ekfac.yaml` +# Run with `bergson examples/compare_wikitext/ekfac.yaml` run_path: runs/compare_wikitext/ekfac steps: @@ -48,12 +48,13 @@ steps: inversion_cfg: damping_factor: 0.1 - - validate: &validate + - 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 @@ -66,22 +67,7 @@ steps: 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 + nproc_per_node: 8 + nnode: 1 - wandb_project: magic - - - validate: - # YAML merge key: copy the contents of &validate - <<: *validate - run_path: runs/compare_wikitext/validate_ekfac_random_exclude_0 - exclude_zero_scores: true + batch_size: 64 diff --git a/examples/compare_wikitext/magic.yaml b/examples/compare_wikitext/magic.yaml index a8300c7b..f339a4cf 100644 --- a/examples/compare_wikitext/magic.yaml +++ b/examples/compare_wikitext/magic.yaml @@ -1,8 +1,11 @@ -# Compute per-query MAGIC attribution scores (one backward per query, 50 -# queries) with leave-k-out validation, then revalidate the same scores -# excluding zero scores. +# 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/magic/compare_wikitext/magic.yaml` +# Run with `bergson examples/compare_wikitext/magic.yaml` run_path: runs/compare_wikitext/magic steps: @@ -10,40 +13,9 @@ steps: run_path: runs/compare_wikitext/random model: gpt2 overwrite: 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: 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 - - # Export the trained model for q3_ekfac.yaml / q3_trackstar.yaml. - save_models: true - wandb_project: magic - - - validate: - run_path: runs/compare_wikitext/validate_magic_random_exclude_0 - model: gpt2 - overwrite: true - - scores: runs/compare_wikitext/random/scores.pt + # 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 @@ -60,7 +32,12 @@ steps: nnode: 1 batch_size: 256 - num_epochs: 2 + # 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 @@ -68,5 +45,9 @@ steps: lr_end: 0.00008 warmup_steps: 0.25 - exclude_zero_scores: true + num_subsets: 100 + subset_fraction: 0.01 + + save_models: true + save_optimizer_state: last wandb_project: magic diff --git a/examples/compare_wikitext/source_adam.yaml b/examples/compare_wikitext/source_adam.yaml new file mode 100644 index 00000000..8ee0906d --- /dev/null +++ b/examples/compare_wikitext/source_adam.yaml @@ -0,0 +1,108 @@ +# 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 + overwrite: 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 index 6373d654..1e8ecc23 100644 --- a/examples/compare_wikitext/trackstar.yaml +++ b/examples/compare_wikitext/trackstar.yaml @@ -1,15 +1,16 @@ -# Compute TrackStar attribution scores for a model generated -# using q3_magic.yaml against test[0:50] (one score column per query), then -# run leave-k-out validation on those scores, with and without zero-score -# exclusion. +# 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/magic/compare_wikitext/trackstar.yaml` +# Run with `bergson examples/compare_wikitext/trackstar.yaml` run_path: runs/compare_wikitext/trackstar steps: - trackstar: index_cfg: - projection_dim: 16 + projection_dim: 32 run_path: runs/compare_wikitext/trackstar model: runs/compare_wikitext/random/retrained/base overwrite: true @@ -37,12 +38,13 @@ steps: score_cfg: batch_size: 1024 - - validate: &validate + - 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 @@ -55,22 +57,7 @@ steps: 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 + nproc_per_node: 8 + nnode: 1 - wandb_project: magic - - - validate: - # YAML merge key: copy the contents of &validate - <<: *validate - run_path: runs/compare_wikitext/validate_trackstar_random_exclude_0 - exclude_zero_scores: true + 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 From ff48977c52c15184c48a019bccf6c18cd9c43c4b Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 05:56:00 +0000 Subject: [PATCH 4/7] compare_wikitext: add plain SOURCE, share the interval run via resume SGD-variant SOURCE scores mean rho 0.221 vs 0.243 for the Adam variant -- the lotus-era SGD advantage does not reproduce on this recipe. Both source yamls now resume the shared interval training run instead of re-training it. --- examples/compare_wikitext/README.md | 1 + examples/compare_wikitext/source.yaml | 102 +++++++++++++++++++++ examples/compare_wikitext/source_adam.yaml | 3 +- 3 files changed, 105 insertions(+), 1 deletion(-) create mode 100644 examples/compare_wikitext/source.yaml diff --git a/examples/compare_wikitext/README.md b/examples/compare_wikitext/README.md index c55d36b9..ed587d71 100644 --- a/examples/compare_wikitext/README.md +++ b/examples/compare_wikitext/README.md @@ -4,4 +4,5 @@ | 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/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 index 8ee0906d..36eaa42a 100644 --- a/examples/compare_wikitext/source_adam.yaml +++ b/examples/compare_wikitext/source_adam.yaml @@ -18,7 +18,8 @@ steps: - train: run_path: runs/compare_wikitext/interval model: gpt2 - overwrite: true + # Shared with source.yaml; resumes instantly when the run exists. + resume: true data: dataset: EleutherAI/bergson-wikitext-512-chunks From 0dfd75b316ab2084ae2891f966ccc7476e33bc66 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 23:29:39 +0900 Subject: [PATCH 5/7] Update docs --- bergson/magic/config.py | 6 +++--- docs/{pipeline.rst => building_blocks.rst} | 4 ++-- docs/magic.rst | 8 ++++---- examples/compare_wikitext/magic.yaml | 1 + examples/magic/gpt2_wikitext_replication.yaml | 3 ++- 5 files changed, 12 insertions(+), 10 deletions(-) rename docs/{pipeline.rst => building_blocks.rst} (99%) 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/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/magic.rst b/docs/magic.rst index 1de8aa6d..24faf73a 100644 --- a/docs/magic.rst +++ b/docs/magic.rst @@ -76,10 +76,10 @@ After a run completes, ``run_cfg.run_path`` contains: * ``validation.csv`` — leave-subset-out validation results (if validation was run). -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 +89,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/magic.yaml b/examples/compare_wikitext/magic.yaml index f339a4cf..874f32b7 100644 --- a/examples/compare_wikitext/magic.yaml +++ b/examples/compare_wikitext/magic.yaml @@ -45,6 +45,7 @@ steps: lr_end: 0.00008 warmup_steps: 0.25 + skip_validation: false num_subsets: 100 subset_fraction: 0.01 diff --git a/examples/magic/gpt2_wikitext_replication.yaml b/examples/magic/gpt2_wikitext_replication.yaml index bf39eccf..8e85885d 100644 --- a/examples/magic/gpt2_wikitext_replication.yaml +++ b/examples/magic/gpt2_wikitext_replication.yaml @@ -46,7 +46,8 @@ steps: lr_start: 1e-6 lr_end: 0.00008 warmup_steps: 0.25 - + + skip_validation: false num_subsets: 100 subset_fraction: 0.01 From 7e77804e75d80bf77e04c780be5943ae2f27b311 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 23:40:34 +0900 Subject: [PATCH 6/7] Update docs --- docs/magic.rst | 43 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 35 insertions(+), 8 deletions(-) diff --git a/docs/magic.rst b/docs/magic.rst index 24faf73a..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,9 +90,18 @@ 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.) Smoothness ---------- From c2f0a0f4e8dd90087181d9b1917ff625fd84b174 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:42:17 +0000 Subject: [PATCH 7/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/magic/gpt2_wikitext_replication.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/magic/gpt2_wikitext_replication.yaml b/examples/magic/gpt2_wikitext_replication.yaml index 8e85885d..2cf11542 100644 --- a/examples/magic/gpt2_wikitext_replication.yaml +++ b/examples/magic/gpt2_wikitext_replication.yaml @@ -46,7 +46,7 @@ steps: lr_start: 1e-6 lr_end: 0.00008 warmup_steps: 0.25 - + skip_validation: false num_subsets: 100 subset_fraction: 0.01